diff --git a/pyproject.toml b/pyproject.toml index c7776b3..1dc858f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ authors = [ ] dependencies = [ "wslink>=2.5.8,<3", - "trame-common>=1.2.7", + "trame-common>=1.3.0", "more-itertools", ] requires-python = ">=3.10" @@ -73,6 +73,7 @@ xfail_strict = true minversion = "6.0" addopts = ["-ra", "--strict-config", "--strict-markers"] filterwarnings = ["error"] +asyncio_default_fixture_loop_scope = "function" testpaths = [ "tests", ] diff --git a/src/trame_server/core.py b/src/trame_server/core.py index 529188e..d66b21a 100644 --- a/src/trame_server/core.py +++ b/src/trame_server/core.py @@ -1,5 +1,6 @@ from __future__ import annotations +import argparse import asyncio import inspect import logging @@ -31,6 +32,42 @@ def set_default_client_type(value: ClientType) -> None: DEFAULT_CLIENT_TYPE = value +class _MCPTransports(argparse.Action): + """Parse ``--mcp-transport``, which is comma-separated and repeatable. Each name + is checked as it is parsed, so that a wrong one fails like any other wrong + value.""" + + def __call__(self, parser, namespace, values, option_string=None): + from .mcp.core import TRANSPORTS # noqa: PLC0415 + + names = [*(getattr(namespace, self.dest, None) or ())] + names += [n.strip() for n in values.split(",") if n.strip()] + unknown = [n for n in names if n not in (*TRANSPORTS, "none")] + if unknown: + parser.error( + f"{option_string}: unknown transport {unknown[0]!r} " + f"(choose from {', '.join(TRANSPORTS)}, or none)" + ) + if "none" in names and len(set(names)) > 1: + parser.error(f"{option_string}: none can't be combined with a transport") + if getattr(namespace, "mcp_no_transport", False) and names != ["none"]: + parser.error(f"{option_string}: --no-mcp-transport serves none") + setattr(namespace, self.dest, tuple(dict.fromkeys(names))) + + +class _NoMCPTransport(argparse.Action): + """Parse ``--no-mcp-transport``, which contradicts a ``--mcp-transport`` that + names a transport, whichever comes first.""" + + def __init__(self, option_strings, dest, **kwargs): + super().__init__(option_strings, dest, nargs=0, **kwargs) + + def __call__(self, parser, namespace, _values, option_string=None): + if (getattr(namespace, "mcp_transport", None) or ("none",)) != ("none",): + parser.error(f"{option_string}: --mcp-transport names a transport") + setattr(namespace, self.dest, True) + + class Server: """ Server implementation for trame. @@ -80,6 +117,13 @@ def __init__( self._root_protocol = None self._protocols_to_configure = [] + # Agents (MCP): created on first access to self.mcp, root only + self._mcp = None + self._child_mcp = None # a child's view of its root's + self._mcp_bound = False + self._mcp_endpoint_mounted = False + self._server_options = None + # ENV variable mapping settings self.hot_reload = "--hot-reload" in sys.argv or bool( os.getenv("TRAME_HOT_RELOAD", None) @@ -103,7 +147,10 @@ def __init__( # Shared state + reserve internal keys if parent_server is None: self._state = State( - self.translator, commit_fn=self._push_state, hot_reload=self.hot_reload + self.translator, + commit_fn=self._push_state, + hot_reload=self.hot_reload, + get_mcp=self._get_mcp, ) for key in [ "scripts", @@ -124,6 +171,7 @@ def __init__( internal=parent_server._state, commit_fn=self._push_state, hot_reload=self.hot_reload, + get_mcp=self._get_mcp, ) # Controller @@ -420,6 +468,59 @@ def cli(self): action="store_true", ) + # Agents (MCP) + self._cli_parser.add_argument( + "--mcp", + dest="mcp", + action=argparse.BooleanOptionalAction, + default=None, + help="""Turn agent access through MCP on or off. --no-mcp holds for the + whole run: no transport is served and agents in the process + are refused too.""", + ) + self._cli_parser.add_argument( + "--mcp-transport", + dest="mcp_transport", + action=_MCPTransports, + metavar="TRANSPORT", + default=None, + help="""How agents outside the process connect at start, instead of + the application's choice: http, or none for agents in the + process only. Comma-separated, repeatable. The application + may change it once running.""", + ) + self._cli_parser.add_argument( + "--no-mcp-transport", + dest="mcp_no_transport", + action=_NoMCPTransport, + default=False, + help="""Serve no transport for the whole run, so that only agents in + the process can connect, when access is on.""", + ) + self._cli_parser.add_argument( + "--mcp-http-token", + dest="mcp_http_token", + metavar="TOKEN", + help="Bearer token required on the HTTP transport (or set TRAME_MCP_HTTP_TOKEN)", + ) + self._cli_parser.add_argument( + "--mcp-log", + dest="mcp_log", + action="store_true", + help="""Log what agents call, at INFO on the trame_server.mcp loggers + (or set TRAME_MCP_LOG)""", + ) + self._cli_parser.add_argument( + "--mcp-deny", + dest="mcp_deny", + action="append", + metavar="PERMISSION", + default=None, + help="""Deny agents these permissions for the whole run: the application + cannot allow them again. Comma-separated, repeatable (or set + TRAME_MCP_DENY)""", + ) + CoreServer.add_arguments(self._cli_parser) return self._cli_parser @@ -442,6 +543,58 @@ def context(self) -> State: """ return self._context + @property + def mcp(self): + """ + The MCP server agents connect to, created on first access. A child server + shares its root's, and the names it exposes objects under go through its + translator. + + :return: The server MCP + :rtype: trame_server.mcp.MCP + """ + if self.root_server != self: + if self._child_mcp is None: + from .mcp.core import ChildMCP # noqa: PLC0415 + + self._child_mcp = ChildMCP(self.root_server.mcp, self.translator) + return self._child_mcp + + if self._mcp is None: + from .mcp import MCP # noqa: PLC0415 + + self._mcp = MCP(self) + return self._mcp + + def _get_mcp(self): + return self.mcp + + def _bind_mcp(self, wslink_server): + """Called once the web server is bound. Mount the MCP endpoint, whether or not + self.mcp exists yet or serves http now, since routes freeze at start and the + app may add the transport later. Only --no-mcp and --no-mcp-transport leave it + out.""" + from .mcp.core import transports_withheld # noqa: PLC0415 + + self._mcp_bound = True + if not transports_withheld(self._server_options): + from .mcp.endpoint import mount # noqa: PLC0415 + + self._mcp_endpoint_mounted = mount(self, wslink_server) + if self._mcp is not None: + self._mcp._start() + + def _mcp_ready(self): + """Called when the server is ready: announce agent access when it is live.""" + if self._mcp is not None: + self._mcp._on_ready() + if self._mcp._exposes_anything(): + return + if getattr(self._server_options, "mcp", None): + logging.getLogger("trame_server.mcp").warning( + "--mcp is given, but this app exposes nothing to agents" + ) + @property def controller(self) -> Controller: """ diff --git a/src/trame_server/mcp/__init__.py b/src/trame_server/mcp/__init__.py new file mode 100644 index 0000000..413c847 --- /dev/null +++ b/src/trame_server/mcp/__init__.py @@ -0,0 +1,32 @@ +"""Agent access to trame applications through MCP (Model Context Protocol). + +``server.mcp`` is the MCP server of a trame server: it serves the tools, +resources and prompts an application declares, and the objects it exposes, +to agents in other processes through its transports (HTTP, on the same port +as the application) and to agents inside it (``server.mcp.client()``). See +:class:`trame_server.mcp.core.MCP`. +""" + +from trame_common.obj.mcp import Progress + +from trame_server.mcp.client import MCPClient, MCPError +from trame_server.mcp.content import Image +from trame_server.mcp.core import MCP, MCPDisabledError +from trame_server.mcp.endpoint import HTTPTransport, claude_command, codex_command +from trame_server.mcp.permissions import PermissionDeniedError +from trame_server.mcp.schema import ToolError, ValidationError + +__all__ = [ + "MCP", + "HTTPTransport", + "Image", + "MCPClient", + "MCPDisabledError", + "MCPError", + "PermissionDeniedError", + "Progress", + "ToolError", + "ValidationError", + "claude_command", + "codex_command", +] diff --git a/src/trame_server/mcp/budget.py b/src/trame_server/mcp/budget.py new file mode 100644 index 0000000..4e05d5e --- /dev/null +++ b/src/trame_server/mcp/budget.py @@ -0,0 +1,136 @@ +"""Keeping replies small: values nobody asked for are capped, and every +reply stays within ``mcp.max_response_tokens``. + +Whatever is cut short says so in a ``more`` entry, which names the call that +returns the rest when there is one. +""" + +from __future__ import annotations + +import json + +__all__ = ["CHARS_PER_TOKEN", "MAX_ITEMS", "MAX_STRING", "cap", "shrink", "size"] + +# A rough estimate, good enough to stay under clients' limits +CHARS_PER_TOKEN = 4 +# Unrequested values: longer lists and strings are cut +MAX_ITEMS = 20 +MAX_STRING = 400 +# Below this, a string is not worth cutting to fit the budget +MIN_STRING = 40 + + +def size(value) -> int: + """The length of ``value`` as the compact JSON text a reply carries.""" + return len( + json.dumps(value, separators=(",", ":"), ensure_ascii=False, default=str) + ) + + +def _cut_string(text, length): + return text[: max(1, length - 1)] + "…" + + +def cap(value): + """``(value, note)``: ``value`` with its lists cut to ``MAX_ITEMS`` and + its strings to ``MAX_STRING`` characters, at any depth. ``note`` says + what was cut (None when nothing was).""" + cut = [] + + def visit(v, top): + if isinstance(v, str): + if len(v) > MAX_STRING: + cut.append(f"{MAX_STRING} of {len(v)} characters shown" if top else "") + return _cut_string(v, MAX_STRING) + return v + if isinstance(v, list): + items = v + if len(v) > MAX_ITEMS: + cut.append(f"{MAX_ITEMS} of {len(v)} items shown" if top else "") + items = v[:MAX_ITEMS] + return [visit(i, False) for i in items] + if isinstance(v, dict): + return {k: visit(i, False) for k, i in v.items()} + return v + + result = visit(value, True) + if not cut: + return result, None + notes = [c for c in cut if c] + return result, notes[0] if notes else "shortened" + + +def _cuttable(v): + return (isinstance(v, list) and len(v) > 1) or ( + isinstance(v, str) and len(v) > MIN_STRING + ) + + +def _largest(value): + """The path to the node to cut next, or None. Lists and strings are cut, and + dicts are walked into, through their largest member each time.""" + path, node = [], value + while True: + if isinstance(node, dict): + if not node: + return None + key = max(node, key=lambda k: size(node[k])) + if not (_cuttable(node[key]) or isinstance(node[key], (dict, list))): + return None + path.append(key) + node = node[key] + continue + if isinstance(node, list): + if not node: + return None + sizes = [size(v) for v in node] + big = max(range(len(node)), key=sizes.__getitem__) + # One item holding most of the list: cut inside it instead + inner = node[big] + if ( + sizes[big] * 2 > sum(sizes) + and isinstance(inner, (dict, list, str)) + and (len(node) == 1 or _cuttable(inner)) + ): + path.append(big) + node = inner + continue + return path if len(node) > 1 else None + if isinstance(node, str) and len(node) > MIN_STRING: + return path + return None + + +def shrink(value, budget, totals=None): + """``(value, cuts)``: ``value`` (a dict) with its largest lists and + strings cut until its JSON fits in ``budget`` characters, or nothing is + left to cut. ``cuts`` maps each top-level key cut into to a note. + ``totals`` gives the full length of top-level values already cut (by + ``cap``), for the notes.""" + value = json.loads(json.dumps(value, default=str)) # a copy we may cut + totals, lengths, cuts = totals or {}, {}, {} + while size(value) > budget: + path = _largest(value) + if not path: + break + parent = value + for key in path[:-1]: + parent = parent[key] + node = parent[path[-1]] + top = path[0] + if len(path) == 1: + lengths.setdefault(top, totals.get(top, len(node))) + if isinstance(node, list): + parent[path[-1]] = node[: max(1, len(node) // 2)] + else: + parent[path[-1]] = _cut_string(node, len(node) // 2) + cuts[top] = None + for top in cuts: + node = value[top] + if top in lengths and isinstance(node, list): + cuts[top] = f"{len(node)} of {lengths[top]} items shown" + elif top in lengths and isinstance(node, str): + cuts[top] = f"{len(node) - 1} of {lengths[top]} characters shown" + else: + cuts[top] = "shortened" + return value, cuts diff --git a/src/trame_server/mcp/call_log.py b/src/trame_server/mcp/call_log.py new file mode 100644 index 0000000..f307a0a --- /dev/null +++ b/src/trame_server/mcp/call_log.py @@ -0,0 +1,120 @@ +"""The log of what agents call (``--mcp-log``). + +Every tool call, resource read and prompt is logged once, when it ends, at +INFO on this module's logger: what was called, with its arguments +shortened, how it ended and how long it took. ``--mcp-log`` or +``TRAME_MCP_LOG`` set the ``trame_server.mcp`` loggers to INFO; where the +lines go is up to the application's logging configuration, as for any +other library. + +``MCPProtocol.handle`` records who is asking in a context variable, which +follows the request into the task that runs it. The caller is the start of +the HTTP session id when there is one, else the name the client gives +(``clientInfo``), which is how in-process clients and clients of MCP +2026-07-28 are named. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import contextvars +import json +import logging +import time + +from trame_server.mcp.permissions import PermissionDeniedError + +__all__ = ["logged", "request", "shorten"] + +logger = logging.getLogger(__name__) + +# How much of an argument a line shows +MAX_STRING = 60 +MAX_ITEMS = 8 +MAX_DEPTH = 3 +MAX_WHAT = 300 +MAX_ERROR = 200 + +_session = contextvars.ContextVar("trame_mcp_session", default="-") + + +@contextlib.contextmanager +def request(caller): + """Mark the block as handling a request from ``caller`` (None when unknown), + so that the calls made in it are logged under that name.""" + token = _session.set(caller or "-") + try: + yield + finally: + _session.reset(token) + + +def shorten(value, depth=0) -> str: + """``value`` as compact JSON-like text: bare keys, long strings replaced + by their length, long lists and dicts cut, deep nesting elided.""" + if isinstance(value, str): + if len(value) > MAX_STRING: + return f"" + return json.dumps(value) + if isinstance(value, dict): + if depth >= MAX_DEPTH: + return "{...}" + items = [ + f"{k}: {shorten(v, depth + 1)}" for k, v in list(value.items())[:MAX_ITEMS] + ] + if len(value) > MAX_ITEMS: + items.append(f"+{len(value) - MAX_ITEMS}") + return "{" + ", ".join(items) + "}" + if isinstance(value, (list, tuple)): + if depth >= MAX_DEPTH: + return "[...]" + items = [shorten(v, depth + 1) for v in value[:MAX_ITEMS]] + if len(value) > MAX_ITEMS: + items.append(f"+{len(value) - MAX_ITEMS}") + return "[" + ", ".join(items) + "]" + if value is None or isinstance(value, (bool, int, float)): + return json.dumps(value) + return f"<{type(value).__name__}>" + + +def _describe_error(error) -> str: + text = str(error).strip() + first = text.splitlines()[0] if text else "" + if not isinstance(error, (ValueError, LookupError, PermissionError)): + first = f"{type(error).__name__}: {first}" if first else type(error).__name__ + if len(first) > MAX_ERROR: + first = first[: MAX_ERROR - 1] + "…" + return first + + +@contextlib.contextmanager +def logged(item_type, name, arguments=None): + """Log the call run in the block when it ends. ``item_type`` is "tool", + "resource" or "prompt".""" + if not logger.isEnabledFor(logging.INFO): + yield + return + started = time.monotonic() + outcome, error = "ok", None + try: + yield + except asyncio.CancelledError: + outcome = "cancelled" + raise + except PermissionDeniedError as e: + outcome, error = "denied", e + raise + except BaseException as e: + outcome, error = "error", e + raise + finally: + what = name if item_type == "tool" else f"{item_type} {name}" + if arguments: + what = f"{what} {shorten(arguments)}" + if len(what) > MAX_WHAT: + what = what[: MAX_WHAT - 1] + "…" + text = f"{what} -> {outcome} {time.monotonic() - started:.2f}s" + if error is not None: + text += f": {_describe_error(error)}" + logger.info("mcp[%s] %s", _session.get(), text) diff --git a/src/trame_server/mcp/client.py b/src/trame_server/mcp/client.py new file mode 100644 index 0000000..8579699 --- /dev/null +++ b/src/trame_server/mcp/client.py @@ -0,0 +1,143 @@ +"""Agents inside the application's process: ``server.mcp.client()``. + +An in-process client needs no transport. It hands the protocol the +requests an agent connected over a transport would send (MCP 2026-07-28), +so what it sees and may do is the same: ``mcp.enabled``, permissions, +argument validation, calls one at a time and the call log all apply. +""" + +from __future__ import annotations + +import inspect +import itertools +import logging + +from trame_server.mcp.protocol import ( + CAPABILITIES_KEY, + CLIENT_INFO_KEY, + MODERN_VERSIONS, + VERSION_KEY, +) + +logger = logging.getLogger("trame_server.mcp") + +__all__ = ["MCPClient", "MCPError"] + +# What a 2026-07-28 result carries for transports, not for the caller +ENVELOPE = ("resultType", "_meta", "ttlMs", "cacheScope") + + +class MCPError(Exception): + """Raised for a request the server cannot serve at all, such as a call to an + unknown tool. It carries the JSON-RPC error's ``code``, ``message`` and + ``data``. A tool that fails still returns a result, with ``isError`` set.""" + + def __init__(self, code, message, data=None): + super().__init__(message) + self.code = code + self.message = message + self.data = data + + +class MCPClient: + """A client for an agent inside the application's process, returned by + ``server.mcp.client()``. It lists what agents see and makes the calls + they make. + + Results are the dicts MCP defines: a tool's result holds its ``content``, + and ``isError`` when the tool failed, which an agent loop hands back to + its model. Requests the server cannot serve raise :class:`MCPError`. + Cancelling the task that awaits a call cancels the call. + """ + + def __init__(self, mcp, name="in-process"): + self._mcp = mcp + self._name = name + self._ids = itertools.count(1) + + @property + def name(self): + """The name this client gives the server, shown in the call log.""" + return self._name + + @property + def instructions(self): + """What agents receive when they connect: the application's text, + then the generated overview of what it exposes.""" + return self._mcp._agent_instructions() + + async def list_tools(self) -> list[dict]: + """The tools agents may call, each with its ``name``, + ``description``, ``inputSchema`` and ``annotations``.""" + return (await self._request("tools/list"))["tools"] + + async def call_tool(self, name, arguments=None, *, on_progress=None) -> dict: + """Call a tool and return its result: ``content``, and ``isError`` when the + tool failed. + + With ``on_progress``, the call asks for the progress that its tool or + method reports, and each report calls ``on_progress(value, total, + message)``, which may be async. These calls run on the event loop, before + the call returns. + """ + params = {"name": name, "arguments": arguments or {}} + if on_progress is None: + return await self._request("tools/call", params) + + async def notify(notification): + fields = notification["params"] + try: + done = on_progress( + fields["progress"], fields.get("total"), fields.get("message") + ) + if inspect.isawaitable(done): + await done + except Exception: # the call goes on without it + logger.exception("on_progress failed") + + return await self._request("tools/call", params, notify=notify) + + async def list_resources(self) -> list[dict]: + """The resources agents may read, each with its ``uri`` and + ``name``.""" + return (await self._request("resources/list"))["resources"] + + async def read_resource(self, uri) -> list[dict]: + """Read a resource. Returns its contents, each with its ``uri``, + ``mimeType`` and ``text`` (or ``blob``, base64-encoded bytes).""" + return (await self._request("resources/read", {"uri": uri}))["contents"] + + async def list_prompts(self) -> list[dict]: + """The prompts agents may get, each with its ``name`` and + ``arguments``.""" + return (await self._request("prompts/list"))["prompts"] + + async def get_prompt(self, name, arguments=None) -> dict: + """Get a prompt. Returns its ``messages`` (and ``description``).""" + return await self._request( + "prompts/get", {"name": name, "arguments": arguments or {}} + ) + + async def _request(self, method, params=None, notify=None): + from trame_server import __version__ # noqa: PLC0415 + + rpc_id = next(self._ids) + meta = { + VERSION_KEY: MODERN_VERSIONS[0], + CAPABILITIES_KEY: {}, + CLIENT_INFO_KEY: {"name": self._name, "version": __version__}, + } + if notify is not None: + meta["progressToken"] = rpc_id # unique among its running requests + message = { + "jsonrpc": "2.0", + "id": rpc_id, + "method": method, + "params": {**(params or {}), "_meta": meta}, + } + # The client itself scopes its request ids, as a session would + response = await self._mcp.protocol.handle(message, session=self, notify=notify) + if "error" in response: + error = response["error"] + raise MCPError(error["code"], error["message"], error.get("data")) + return {k: v for k, v in response["result"].items() if k not in ENVELOPE} diff --git a/src/trame_server/mcp/content.py b/src/trame_server/mcp/content.py new file mode 100644 index 0000000..290a9c3 --- /dev/null +++ b/src/trame_server/mcp/content.py @@ -0,0 +1,90 @@ +"""Coercion of Python return values into MCP tool-result content.""" + +from __future__ import annotations + +import base64 +import json +from pathlib import Path + +__all__ = [ + "Image", + "to_tool_result", +] + + +class Image: + """An image that a tool returns. + + .. code-block:: python + + @server.mcp.tool + def screenshot() -> Image: + return Image(png_bytes, mime_type="image/png") + """ + + def __init__(self, data: bytes, mime_type: str = "image/png"): + self.data = data + self.mime_type = mime_type + + @classmethod + def from_file(cls, path, mime_type=None): + path = Path(path) + if mime_type is None: + ext = path.suffix.lower().lstrip(".") + mime_type = f"image/{'jpeg' if ext in ('jpg', 'jpeg') else ext or 'png'}" + return cls(path.read_bytes(), mime_type=mime_type) + + def to_content(self): + return { + "type": "image", + "data": base64.b64encode(self.data).decode("ascii"), + "mimeType": self.mime_type, + } + + +def _default_encoder(value): + if isinstance(value, Image): + return f"" + return repr(value) + + +def to_tool_result(value, structured=False) -> dict: + """Convert a tool's return value into an MCP ``tools/call`` result. + + Anything but text and images becomes compact JSON text. That includes None, + which becomes ``null`` rather than an empty reply that would leave the agent + guessing. With ``structured``, a dict is also sent as ``structuredContent``. + Clients that show both then put the same value in the agent's context twice + (see ``MCP.structured_content``).""" + if isinstance(value, Image): + return {"content": [value.to_content()]} + + if isinstance(value, str): + return {"content": [{"type": "text", "text": value}]} + + if ( + isinstance(value, (list, tuple)) + and value + and all(isinstance(v, (Image, str)) for v in value) + ): + return { + "content": [ + v.to_content() if isinstance(v, Image) else {"type": "text", "text": v} + for v in value + ] + } + + text = json.dumps( + value, default=_default_encoder, separators=(",", ":"), ensure_ascii=False + ) + result = {"content": [{"type": "text", "text": text}]} + if structured and isinstance(value, dict): + result["structuredContent"] = value + return result + + +def to_error_result(message: str) -> dict: + return { + "content": [{"type": "text", "text": message}], + "isError": True, + } diff --git a/src/trame_server/mcp/core.py b/src/trame_server/mcp/core.py new file mode 100644 index 0000000..0bbef4b --- /dev/null +++ b/src/trame_server/mcp/core.py @@ -0,0 +1,803 @@ +"""The MCP server of a trame server: ``server.mcp``.""" + +from __future__ import annotations + +import inspect +import logging +import os +import re +from dataclasses import dataclass, field + +from trame_common.obj.mcp import permission_names + +from trame_server.mcp.client import MCPClient +from trame_server.mcp.endpoint import HTTPTransport +from trame_server.mcp.objects import OBJECT_TOOLS, ObjectLayer +from trame_server.mcp.permissions import DENY_ENV, Permissions, launch_names +from trame_server.mcp.protocol import MCPProtocol +from trame_server.mcp.schema import ( + check_progress, + docstring, + function_input_schema, + parameter_hints, + type_to_schema, +) +from trame_server.mcp.typeinfo import tool_reference_schema + +logger = logging.getLogger("trame_server.mcp") + +__all__ = ["MCP", "MCPDisabledError"] + +HTTP_TOKEN_ENV = "TRAME_MCP_HTTP_TOKEN" +LOG_ENV = "TRAME_MCP_LOG" + +# How agents outside the process can connect (mcp.transports) +TRANSPORTS = ("http",) +DEFAULT_TRANSPORTS = ("http",) + +# Claude Code rejects a tool whose full name mcp____ is longer +MAX_FULL_TOOL_NAME = 64 +# What the MCP spec allows in a tool name +TOOL_NAME = re.compile(r"^[A-Za-z0-9_.-]{1,128}$") + +LIVE_APP_INSTRUCTIONS = ( + "A live trame app, shared with its users: they see your changes at once." +) + + +class MCPDisabledError(PermissionError): + """Raised for a call made while ``mcp.enabled`` is False.""" + + +@dataclass +class ToolEntry: + name: str + description: str + fn: object + input_schema: dict | None = None # None: derived from fn's signature + annotations: dict | None = None + concurrent: bool = False # runs without the call lock + requires: tuple = () # permission names (mcp.permissions) + cached_schema: tuple = field(default=(None, None), repr=False) + + +@dataclass +class ResourceEntry: + uri: str + name: str + description: str | None + mime_type: str | None + fn: object + + +@dataclass +class PromptEntry: + name: str + description: str | None + arguments: list + fn: object + + +def _tool_annotations(read_only, annotations): + result = dict(annotations or {}) + if read_only: + result["readOnlyHint"] = True + return result or None + + +def _check_transports(names): + unknown = [n for n in names if n not in TRANSPORTS] + if unknown: + msg = ( + f"Unknown MCP transport {unknown[0]!r}; available: {', '.join(TRANSPORTS)} " + f"(an empty list serves agents in the process only)" + ) + raise ValueError(msg) + return tuple(dict.fromkeys(names)) + + +def transports_withheld(args): + """True when the command line serves no transport for the whole run, with + ``--no-mcp`` or ``--no-mcp-transport``.""" + return getattr(args, "mcp", None) is False or bool( + getattr(args, "mcp_no_transport", False) + ) + + +def launch_transports(args): + """The transports the command line chooses: an empty tuple when it withholds + them all or names ``none``, the names given to ``--mcp-transport``, or None + when it leaves the choice to the application.""" + if transports_withheld(args): + return () + names = getattr(args, "mcp_transport", None) + if names is None: + return None + return () if names == ("none",) else names + + +class MCP: + """The MCP server of a trame server, created on first access to + ``server.mcp``. There is one per root server. Child servers share it, and + the names they expose objects under go through their translator. + + It serves what the application declares: first-class tools, resources + and prompts (``@server.mcp.tool``, ``server.mcp.resource(uri)``, + ``@server.mcp.prompt``), and objects through the six ``object_*`` tools. + Objects are the components and models whose class declares something, + reached from the names the application exposes, and the ``state`` object. + Nothing is served until ``enabled`` is set, in code or with ``--mcp`` at + launch. + + Agents outside the process connect through ``transports``: HTTP by + default, on the trame web server's port, with its options on ``http``. + Agents inside the process, such as an assistant built into the + application, use ``client()``. + + Every option is a settable attribute that can change at any time, and + ``configure()`` sets several at once. + """ + + def __init__(self, server): + self._server = server.root_server + self._started = False + self._ready = False + self._name = None + self._enabled = False + self._transports = DEFAULT_TRANSPORTS + self._http = HTTPTransport(self) + self._concurrent_calls = False + self._structured_content = False + self._max_response_tokens = 8000 + self._permissions = Permissions() + #: Text sent to agents when they connect, ahead of the overview trame + #: generates from what the application exposes (default None). + self.instructions = None + + self._tools: dict[str, ToolEntry] = {} + self._resources: dict[str, ResourceEntry] = {} + self._prompts: dict[str, PromptEntry] = {} + + self._objects = ObjectLayer(self) + self._protocol = MCPProtocol(self) + + self._apply_launch_options() + if self._server._mcp_bound: # created once the server was already up + self._start() + if self._server.running: + self._on_ready() + + # ------------------------------------------------------------------------- + # Life cycle (driven by the trame server) + # ------------------------------------------------------------------------- + + def _start(self): + """Called when the server binds its port. From then on, launch options that + turn something on or set a value no longer override assignments. Those + that withhold something (``--no-mcp``, ``--no-mcp-transport`` and + ``--mcp-deny``) still hold.""" + self._started = True + self._permissions._start() + self._objects.start() + for entry in self._tools.values(): + unknown = [n for n in entry.requires if n not in self._permissions] + if unknown: + logger.warning( + "Tool %s requires undeclared permissions %s; it is never offered", + entry.name, + unknown, + ) + self._check_length(entry.name, warn=True) + if self._server._mcp_endpoint_mounted: + self._http._start() + + def _on_ready(self): + self._ready = True + self._objects.audit() + self._http._on_ready() + + # ------------------------------------------------------------------------- + # Launch options + # ------------------------------------------------------------------------- + + def _apply_launch_options(self): + """Apply the command line (``Server.cli``) and the environment: ``enabled``, + ``transports``, the HTTP ``token``, the call log and the denied + permissions. + + They are the choice of the person running the application, so until the + server starts they take precedence over what the application assigns, + which only sets defaults. Those that withhold something (``--no-mcp``, + ``--no-mcp-transport`` and ``--mcp-deny``) hold for the whole run. + """ + args, _ = self._server.cli.parse_known_args() + self._launch_options = { + "enabled": getattr(args, "mcp", None), + "transports": launch_transports(args), + "transports_withheld": transports_withheld(args), + "token": getattr(args, "mcp_http_token", None) + or os.environ.get(HTTP_TOKEN_ENV) + or None, + } + if self._launch_options["enabled"] is not None: + self._enabled = self._launch_options["enabled"] + if self._launch_options["transports"] is not None: + self._transports = self._launch_options["transports"] + if self._launch_options["token"] is not None: + self._http._token = str(self._launch_options["token"]) + if getattr(args, "mcp_log", False) or os.getenv(LOG_ENV): + # Like wslink's --debug, this only sets a level: the handlers are + # the application's + if logger.getEffectiveLevel() > logging.INFO: + logger.setLevel(logging.INFO) + # The flag and the variable both apply, since each withholds something + self._permissions._deny_at_launch( + launch_names( + [*(getattr(args, "mcp_deny", None) or ()), os.environ.get(DENY_ENV)] + ) + ) + + def _launch_overrides(self, option): + """True when a launch option decides ``option`` and the server has not + started yet: the application's assignment is then only a default.""" + if self._started or self._launch_options.get(option) is None: + return False + logger.info( + "MCP %s is set at launch; the application's value is ignored", option + ) + return True + + # ------------------------------------------------------------------------- + # Options + # ------------------------------------------------------------------------- + + def configure( + self, + *, + name=None, + enabled=None, + transports=None, + concurrent_calls=None, + structured_content=None, + instructions=None, + max_response_tokens=None, + ): + """Set several options at once; an option left as None is unchanged. This + is the same as assigning the attributes one by one. The HTTP transport's + options are set with ``http.configure()``.""" + options = { + "name": name, + "enabled": enabled, + "transports": transports, + "concurrent_calls": concurrent_calls, + "structured_content": structured_content, + "instructions": instructions, + "max_response_tokens": max_response_tokens, + } + for key, value in options.items(): + if value is not None: + setattr(self, key, value) + + @property + def server(self): + """The trame server this MCP server belongs to (a root server).""" + return self._server + + @property + def protocol(self): + return self._protocol + + @property + def name(self): + """The name reported to clients and used in the connection commands + (default: the trame server's name).""" + return self._name or self._server.name + + @name.setter + def name(self, value): + self._name = value + + @property + def enabled(self): + """Whether agents can use what the application declares (default False). + + This is the one switch for agent access, whichever way agents connect, + and it can change at any time. While it is False, every call is refused + and the HTTP endpoint answers 503. ``--mcp`` at launch takes precedence + over values the application sets before the server starts; once the + server runs, every assignment applies. ``--no-mcp`` keeps access off for + the whole run (:attr:`enabled_locked`). + """ + return self._enabled + + @enabled.setter + def enabled(self, value): + if value and self.enabled_locked: + logger.info("--no-mcp keeps agent access off for this run") + return + if self._launch_overrides("enabled"): + return + self._enabled = bool(value) + if self._enabled: + self._http._announce() + + @property + def enabled_locked(self): + """True when ``--no-mcp`` keeps agent access off for the whole run. Turning + access on is then ignored, so a UI should disable its switch.""" + return self._launch_options["enabled"] is False + + @property + def transports(self): + """How agents outside the process connect (default ``("http",)``): a tuple + of transport names, empty when only agents inside the process may connect. + + It can change at any time; while ``"http"`` is left out, the HTTP + endpoint answers 503. ``--mcp-transport`` at launch takes precedence over + values the application sets before the server starts. + ``--no-mcp-transport`` and ``--no-mcp`` serve no transport for the whole + run (:attr:`transports_locked`). + """ + return self._transports + + @transports.setter + def transports(self, value): + names = _check_transports((value,) if isinstance(value, str) else value) + if names and self.transports_locked: + flag = "--no-mcp" if self.enabled_locked else "--no-mcp-transport" + logger.info("%s keeps agents in other processes out for this run", flag) + return + if self._launch_overrides("transports"): + return + self._transports = names + if "http" in names: + self._http._announce() + + @property + def transports_locked(self): + """True when ``--no-mcp-transport`` or ``--no-mcp`` serve no transport for + the whole run. Adding a transport is then ignored, so a UI should disable + its switch.""" + return self._launch_options["transports_withheld"] + + def client(self, name="in-process"): + """A client for an agent inside the application's process + (:class:`MCPClient`). It lists and calls what agents may use, without a + transport but under the same rules (``enabled``, permissions). ``name`` + identifies it in the call log.""" + return MCPClient(self, name) + + @property + def http(self): + """The HTTP transport's options (:class:`HTTPTransport`): ``path``, + ``token``, ``allowed_origins`` and the read-only ``url``.""" + return self._http + + @property + def permissions(self): + """The permissions: capabilities that the people running the application can + withhold from agents. ``declare()`` each one, then name the permissions + that a tool, method, property or constructor needs with ``requires=``. + ``--mcp-deny`` or ``TRAME_MCP_DENY`` deny permissions for the whole run + (``permissions.locked``).""" + return self._permissions + + @property + def concurrent_calls(self): + """When False (default), agent calls run one at a time, except those + declared with ``concurrent=True``. When True, every call may run + while others are in progress.""" + return self._concurrent_calls + + @concurrent_calls.setter + def concurrent_calls(self, value): + self._concurrent_calls = bool(value) + + @property + def structured_content(self): + """Whether to also send a tool's dict result as ``structuredContent`` + (default False). The result is always sent as compact JSON text too.""" + return self._structured_content + + @structured_content.setter + def structured_content(self, value): + self._structured_content = bool(value) + + @property + def max_response_tokens(self): + """The maximum size of a reply from the ``object_*`` tools, in tokens of + about 4 characters (default 8000). A reply cut to fit says how to get the + rest.""" + return self._max_response_tokens + + @max_response_tokens.setter + def max_response_tokens(self, value): + if not isinstance(value, int) or value < 500: + msg = ( + f"max_response_tokens must be an integer of at least 500, got {value!r}" + ) + raise ValueError(msg) + self._max_response_tokens = value + + def _agent_instructions(self): + """What clients receive as ``instructions``: the application's own + text, then a generated overview.""" + lines = [LIVE_APP_INSTRUCTIONS, *self._objects.instructions()] + denied = self._permissions.denied + if denied: + names = ", ".join( + f"{n} ({self._permissions.description(n)})" for n in denied + ) + lines.append( + f"Turned off in this application: {names}. What needs them is not " + f"offered, and calls to it are refused." + ) + overview = "\n".join(lines) + if self.instructions: + return f"{self.instructions.strip()}\n\n{overview}" + return overview + + def _check_enabled(self): + if not self._enabled: + msg = "MCP access is currently disabled by the application" + raise MCPDisabledError(msg) + self._permissions._check_launch() + + # ------------------------------------------------------------------------- + # First-class tools, resources and prompts + # ------------------------------------------------------------------------- + + def tool( + self, + fn=None, + *, + name=None, + description=None, + read_only=False, + annotations=None, + concurrent=False, + requires=None, + ): + """Make a function, or a bound method, a tool. Use it as a decorator or call + it: ``server.mcp.tool(self.screenshot)``. + + The type hints give the input schema. ``Annotated[T, "description", + {json-schema}]`` adds a description and constraints, and a parameter + typed with a component or model class takes that object's id or name. + Results: a ``str`` is sent as text, an ``Image`` as an image (a list of + them as several), and anything else as compact JSON text (``null`` for + None). + + :param name: the tool name (default: the function's). + :param description: the text agents read (default: the docstring). + :param read_only: True when it changes nothing (``readOnlyHint``). + :param annotations: extra MCP tool annotations. + :param concurrent: True to let it run while other agent calls are in + progress. + :param requires: the permission, or list of permissions, that a call + needs. + """ + + def decorate(f): + self._add_tool( + f, + name=name, + description=description, + annotations=_tool_annotations(read_only, annotations), + concurrent=concurrent, + requires=requires, + ) + return f + + if fn is None: + return decorate + return decorate(fn) + + def resource(self, uri, *, name=None, description=None, mime_type=None): + """Serve what a function returns as the resource ``uri``. This returns a + decorator; for a method, call it as + ``server.mcp.resource(uri)(self.guide)``.""" + + def decorate(f): + self._add_resource( + f, uri=uri, name=name, description=description, mime_type=mime_type + ) + return f + + return decorate + + def prompt(self, fn=None, *, name=None, description=None): + """Make a function, or a bound method, a prompt. Use it as a decorator or + call it.""" + + def decorate(f): + self._add_prompt(f, name=name, description=description) + return f + + if fn is None: + return decorate + return decorate(fn) + + def _check_length(self, tool_name, warn=False): + full = f"mcp__{self.name}__{tool_name}" + if len(full) <= MAX_FULL_TOOL_NAME: + return + msg = ( + f"Tool {tool_name!r}: its full name {full!r} has {len(full)} characters; " + f"clients such as Claude Code reject more than {MAX_FULL_TOOL_NAME}. " + f"Shorten the tool name or mcp.name" + ) + if warn: + logger.warning(msg) + return + raise ValueError(msg) + + def _add_tool( + self, + fn, + name=None, + description=None, + annotations=None, + concurrent=False, + requires=None, + ): + tool_name = name or fn.__name__ + if tool_name in OBJECT_TOOLS: + msg = f"{tool_name!r} is the name of a built-in object tool; choose another" + raise ValueError(msg) + if tool_name in self._tools: + msg = f"A tool named {tool_name!r} is already registered" + raise ValueError(msg) + if not isinstance(tool_name, str) or not TOOL_NAME.match(tool_name): + msg = f"A tool name is 1-128 of A-Z a-z 0-9 _ - (and .), got {tool_name!r}" + raise ValueError(msg) + if "." in tool_name: + logger.warning( + "Tool %r: OpenAI-style function calling rejects dots in names", + tool_name, + ) + self._check_length(tool_name) + check_progress(fn, f"Tool {tool_name!r}") + self._tools[tool_name] = ToolEntry( + name=tool_name, + description=description or docstring(fn) or "", + fn=fn, + annotations=annotations, + concurrent=bool(concurrent), + requires=permission_names(requires), + ) + + def _add_builtin_tool( + self, fn, name, description, annotations=None, concurrent=False + ): + """Add a built-in tool, whose input schema comes from ``fn``'s signature.""" + self._tools[name] = ToolEntry( + name=name, + description=description, + fn=fn, + input_schema=function_input_schema(fn), + annotations=annotations, + concurrent=concurrent, + ) + + def _input_schema(self, entry): + """The tool's input schema. A schema derived from the signature names object + types, so it is computed again whenever they may have changed.""" + if entry.input_schema is not None: + return entry.input_schema + from trame_common.obj.mcp import options_version # noqa: PLC0415 + + version = options_version() + cached_version, schema = entry.cached_schema + if cached_version != version: + schema = function_input_schema(entry.fn, ref=tool_reference_schema) + entry.cached_schema = (version, schema) + return schema + + def _add_resource(self, fn, uri, name=None, description=None, mime_type=None): + if uri in self._resources: + msg = f"A resource with uri {uri!r} is already registered" + raise ValueError(msg) + self._resources[uri] = ResourceEntry( + uri=uri, + name=name or fn.__name__, + description=description or docstring(fn), + mime_type=mime_type, + fn=fn, + ) + + def _add_prompt(self, fn, name=None, description=None): + prompt_name = name or fn.__name__ + if prompt_name in self._prompts: + msg = f"A prompt named {prompt_name!r} is already registered" + raise ValueError(msg) + + hints = parameter_hints(fn) + arguments = [] + for pname, param in inspect.signature(fn).parameters.items(): + if pname == "self" or param.kind in ( + inspect.Parameter.VAR_POSITIONAL, + inspect.Parameter.VAR_KEYWORD, + ): + continue + argument = { + "name": pname, + "required": param.default is inspect.Parameter.empty, + } + # Annotated[str, "description"], as for tool parameters + argument_description = type_to_schema(hints.get(pname)).get("description") + if argument_description: + argument["description"] = argument_description + arguments.append(argument) + self._prompts[prompt_name] = PromptEntry( + name=prompt_name, + description=description or docstring(fn), + arguments=arguments, + fn=fn, + ) + + def list_tools(self): + """The tools offered to agents: those whose permissions are all allowed. + ``get_tool`` finds the other tools too, so that a call to one of them is + refused with the reason.""" + return [t for t in self._tools.values() if self._permissions.allows(t.requires)] + + def get_tool(self, name): + return self._tools.get(name) + + def list_resources(self): + return list(self._resources.values()) + + def get_resource(self, uri): + return self._resources.get(uri) + + def list_prompts(self): + return list(self._prompts.values()) + + def get_prompt(self, name): + return self._prompts.get(name) + + # ------------------------------------------------------------------------- + # Objects + # ------------------------------------------------------------------------- + + def expose_object(self, source, name, *, description=None, read_only=False): + """Expose an object to agents under ``name``. + + :param source: a component or model, or a function that returns one (or + None). trame calls the function on every agent request, so + ``lambda: self.active`` follows reassignment. A name holds a + component weakly and a model strongly. + :param name: must start with a letter or an underscore, and ``state`` is + reserved. A name that already exposes something else must be + withdrawn first, while exposing the same source again updates its + options. Like a state key, the name goes through the server's + translator: on a child server with the prefix ``left_``, agents see + ``left_viewer`` for ``viewer``. + :param description: what this object is for in the application. Agents + read it next to the name; the class already says what the type is. + :param read_only: agents can call only ``read_only`` methods, and write + no property, of what they reach only this way. Reached by id, an + object is read-only too, unless a writable name (or a function's + result under one), or a property of a writable object, also reaches + it. + """ + self._objects.expose(source, self._translate(name), description, read_only) + + def unexpose_object(self, name): + """Withdraw a name given with :meth:`expose_object`.""" + self._objects.unexpose(self._translate(name)) + + def _translate(self, name): + return self._server.translator.translate_key(name) + + def describe(self, cls): + """What agents currently see of ``cls``, as ``object_types`` shows it, or + None when it is hidden. Pass ``"State"`` for the state object.""" + return self._objects.describe_class(cls) + + def _remove_object(self, obj): + """Withdraw every name exposing ``obj``.""" + self._objects.remove(obj) + + def _add_component_tools(self, component): + """Add the tools, resources and prompts that a newly built component marks + (``@mcp.tool``, ``@mcp.resource``, ``@mcp.prompt``): all of them, or none + when one cannot be added.""" + from trame_common.obj.mcp import ( # noqa: PLC0415 + PROMPT_MARKER, + RESOURCE_MARKER, + TOOL_MARKER, + first_class_members, + ) + + cls = type(component) + added = [] + try: + for attr, marker, options in first_class_members(cls): + fn = getattr(component, attr) + if marker == TOOL_MARKER: + key = options["name"] or attr + self._check_free(self._tools, key, "tool", cls) + self._add_tool( + fn, + name=options["name"], + description=options["description"], + annotations=_tool_annotations( + options["read_only"], options["annotations"] + ), + concurrent=options["concurrent"], + requires=options["requires"], + ) + added.append((self._tools, key)) + elif marker == RESOURCE_MARKER: + key = options["uri"] + self._check_free(self._resources, key, "resource", cls) + self._add_resource(fn, **options) + added.append((self._resources, key)) + elif marker == PROMPT_MARKER: + key = options["name"] or attr + self._check_free(self._prompts, key, "prompt", cls) + self._add_prompt(fn, **options) + added.append((self._prompts, key)) + except Exception: + for registry, key in added: + del registry[key] + raise + + @staticmethod + def _check_free(registry, key, item_type, cls): + entry = registry.get(key) + owner = getattr(getattr(entry, "fn", None), "__self__", None) + if entry is not None and isinstance(owner, cls): + msg = ( + f"{cls.__name__} {owner._id} already added the {item_type} {key!r}: " + f"@mcp.tool, @mcp.resource and @mcp.prompt are for components built " + f"once. Unbind the first (_unbind_annotated_methods()) before building " + f"another, or register per-instance tools with server.mcp.tool()" + ) + raise ValueError(msg) + + def _remove_component_tools(self, component): + """Remove what an unbound component added.""" + for registry in (self._tools, self._resources, self._prompts): + for key, entry in list(registry.items()): + if getattr(entry.fn, "__self__", None) is component: + del registry[key] + + def _expose_state(self, key, **options): + self._objects.expose_state(key, **options) + + def _unexpose_state(self, *keys): + self._objects.state.unexpose(*keys) + + def _exposes_anything(self): + return bool( + self._objects.exposes_anything() + or self._resources + or self._prompts + or any(name not in OBJECT_TOOLS for name in self._tools) + ) + + +class ChildMCP: + """What ``server.mcp`` returns on a child server: the root's MCP, except that + the names it exposes objects under go through the child's translator, as + its state keys and controller names do.""" + + def __init__(self, mcp, translator): + object.__setattr__(self, "_mcp", mcp) + object.__setattr__(self, "_translator", translator) + + def __getattr__(self, name): + return getattr(self._mcp, name) + + def __setattr__(self, name, value): + setattr(self._mcp, name, value) + + def _translate(self, name): + return self._translator.translate_key(name) + + # The root's own, translating with this translator + expose_object = MCP.expose_object + unexpose_object = MCP.unexpose_object diff --git a/src/trame_server/mcp/endpoint.py b/src/trame_server/mcp/endpoint.py new file mode 100644 index 0000000..6b10637 --- /dev/null +++ b/src/trame_server/mcp/endpoint.py @@ -0,0 +1,617 @@ +"""The HTTP transport: the endpoint agents connect to (``server.mcp.http``, +``/mcp`` by default), on the trame web server's port. + +:class:`HTTPTransport` holds its options. :class:`MCPEndpoint` handles +requests without depending on a web server: a host hands it an +:class:`HTTPRequest` and sends back the :class:`HTTPResponse`. +:func:`mount` adapts it to aiohttp, trame's default web server. The server +mounts it when it starts, in every app unless launched with ``--no-mcp`` +or ``--no-mcp-transport``. The route exists from the start because +aiohttp freezes its routes then, while agent access and the transport can +be turned on later. Until ``server.mcp`` exists, is enabled and serves +HTTP, the endpoint answers 503. + +Every request goes through, in order: + +1. ``Origin`` check: browsers always send it on cross-site requests, so a + page from another site can never drive the tools (CSRF). +2. Bearer token, when one is set. Without a token, the ``Host`` must be a + loopback address, which defeats DNS rebinding. +3. ``mcp.enabled``, then ``"http"`` in ``mcp.transports``: 503 while + agent access or the transport is off. +4. ``Content-Type: application/json``, so browsers cannot send a request + without a CORS preflight, which this endpoint never grants. + +The transport is streamable HTTP with JSON responses, for clients of both +eras (see :mod:`trame_server.mcp.protocol`): + +- A request naming its protocol version in ``_meta`` (2026-07-28 and + later) is served statelessly. Its ``MCP-Protocol-Version``, + ``Mcp-Method`` and ``Mcp-Name`` headers must match its body, else 400 + (``HeaderMismatch``). An unsupported version or missing metadata gets + 400, and an unknown method 404. No session id is given or read. +- A client that opens with ``initialize`` is served as earlier revisions + say: the response carries an ``Mcp-Session-Id`` that only scopes request + ids for ``notifications/cancelled``. + +Either way, a request whose client disconnects is cancelled. A request +that asks for progress gets its response as Server-Sent Events once its +first notification goes out: the notifications, then the response. +Without a notification, the response stays plain JSON. +""" + +from __future__ import annotations + +import asyncio +import base64 +import contextlib +import hmac +import ipaddress +import json +import logging +import secrets +import shlex +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass, field +from urllib.parse import urlsplit + +from trame_server.mcp.protocol import ( + HEADER_MISMATCH, + INVALID_REQUEST, + METHOD_NOT_FOUND, + MODERN_VERSIONS, + PARSE_ERROR, + check_request, + request_version, +) + +logger = logging.getLogger("trame_server.mcp") + +__all__ = [ + "HTTPRequest", + "HTTPResponse", + "HTTPTransport", + "MCPEndpoint", + "claude_command", + "codex_command", + "is_loopback", + "mount", +] + +WILDCARD_HOSTS = ("", "0.0.0.0", "::") +SERVER_ERROR = -32000 +SESSION_HEADER = "Mcp-Session-Id" +# Seconds between checks that the client of a running request is still there +DISCONNECT_POLL = 0.5 +# The body field each method mirrors into Mcp-Name (2026-07-28) +NAMED_BY = {"tools/call": "name", "prompts/get": "name", "resources/read": "uri"} +# How a header value outside plain ASCII is sent (2026-07-28) +BASE64_PREFIX, BASE64_SUFFIX = "=?base64?", "?=" +# A response sent as Server-Sent Events (X-Accel-Buffering: proxies pass +# each event on at once) +EVENT_STREAM_HEADERS = { + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + "X-Accel-Buffering": "no", +} + + +def is_loopback(host) -> bool: + """True for ``localhost`` names and loopback IP literals.""" + if not host: + return False + host = host.strip("[]").lower() + if host == "localhost" or host.endswith(".localhost"): + return True + try: + return ipaddress.ip_address(host).is_loopback + except ValueError: + return False + + +def claude_command(name, url, token=None): + """A Claude Code command starting a session with this server added for + that session only (``--mcp-config``), next to the user's own MCP + servers.""" + server = {"type": "http", "url": url} + if token: + server["headers"] = {"Authorization": f"Bearer {token}"} + config = json.dumps({"mcpServers": {name: server}}, separators=(",", ":")) + return f"claude --mcp-config {shlex.quote(config)}" + + +def codex_command(name, url, token=None): + """A Codex command starting a session with this server added for that + session only: each ``-c`` overrides one config value (TOML) for the + run, here the server's entry under ``mcp_servers``.""" + overrides = [f"mcp_servers.{name}.url={json.dumps(url)}"] + if token: + header = json.dumps(f"Bearer {token}") + overrides.append(f"mcp_servers.{name}.http_headers={{Authorization={header}}}") + return " ".join(["codex", *(f"-c {shlex.quote(o)}" for o in overrides)]) + + +class HTTPTransport: + """The options of the HTTP transport (``server.mcp.http``), which is served + while ``mcp.transports`` includes ``"http"``, as it does by default. + + Every option is a settable attribute, and ``configure()`` sets several at + once. ``path`` is fixed once the server has started; the other options + can change at any time. + """ + + def __init__(self, mcp): + self._mcp = mcp + self._path = "/mcp" + self._token = None + self._token_generated = False + self._allowed_origins = () + self._announced = False + + def configure(self, *, path=None, token=None, allowed_origins=None): + """Set several options at once; an option left as None is unchanged. This + is the same as assigning the attributes one by one.""" + options = {"path": path, "token": token, "allowed_origins": allowed_origins} + for key, value in options.items(): + if value is not None: + setattr(self, key, value) + + @property + def path(self): + """The endpoint's path on the trame web server. It is fixed once the server + has started.""" + return self._path + + @path.setter + def path(self, value): + if self._mcp._started: + msg = "MCP 'path' cannot change once the server has started" + raise RuntimeError(msg) + self._path = value if value.startswith("/") else f"/{value}" + + @property + def token(self): + """The bearer token the endpoint requires, or None. + + ``--mcp-http-token`` or ``TRAME_MCP_HTTP_TOKEN`` take precedence over values + the application sets before the server starts. When the server listens + beyond loopback and no token is set, one is generated at startup and printed + with the connection instructions. + """ + return self._token + + @token.setter + def token(self, value): + if self._mcp._launch_overrides("token"): + return + self._token = str(value) if value else None + self._token_generated = False + + @property + def allowed_origins(self): + """Extra browser origins (``"https://host:port"``) allowed to call the + endpoint. Loopback origins and the endpoint's own origin are always + allowed; requests without an ``Origin`` header are unaffected.""" + return self._allowed_origins + + @allowed_origins.setter + def allowed_origins(self, value): + if isinstance(value, str): + value = (value,) + self._allowed_origins = tuple(o.rstrip("/") for o in value) + + @property + def url(self): + """The endpoint's URL, as agents on this machine reach it.""" + host = self._bind_host + if not host or host in WILDCARD_HOSTS: + host = "localhost" + elif ":" in host: + host = f"[{host}]" + return f"http://{host}:{self._mcp.server.port}{self._path}" + + @property + def _bind_host(self): + options = getattr(self._mcp.server, "_server_options", None) + return getattr(options, "host", None) + + @property + def _mounted(self): + return self._mcp.server._mcp_endpoint_mounted + + @property + def _serving(self): + return self._mounted and "http" in self._mcp.transports + + @property + def _live(self): + return self._mcp.enabled and self._serving + + def _start(self): + if self._token is None and not is_loopback(self._bind_host): + # Listening beyond this machine: never serve an unauthenticated + # automation API to the network. + self._token = generate_token() + self._token_generated = True + + def _on_ready(self): + if self._live: + self._announce() + elif not self._mcp.enabled and self._mounted: + logger.info( + "%s is mounted, but agent access is off (set mcp.enabled = True or " + "launch with --mcp)", + self.url, + ) + elif self._mounted: + logger.info( + "%s is mounted, but the HTTP transport is off (add it to " + "mcp.transports)", + self.url, + ) + + def _announce(self): + """Print how to connect, the first time access is live.""" + if self._announced or not (self._mcp._ready and self._live): + return + self._announced = True + from trame_server.utils.server import print_mcp_informations # noqa: PLC0415 + + print_mcp_informations(self._mcp.server) + + +@dataclass +class HTTPRequest: + """What the endpoint needs of a request.""" + + method: str + headers: Mapping[str, str] # lower-case names + host: str | None # the host name the client addressed, without the port + authority: str | None # the ``Host`` header as sent: host[:port] + read: Callable[[], Awaitable[bytes]] + disconnected: Callable[[], bool] = lambda: False + # Starts the response as a stream with these headers; returns an async + # function writing text to it. None: the host can't stream responses. + open_stream: ( + Callable[[dict], Awaitable[Callable[[str], Awaitable[None]]]] | None + ) = None + + @property + def content_type(self): + value = self.headers.get("content-type") or "" + return value.split(";")[0].strip().lower() + + +@dataclass +class HTTPResponse: + status: int + body: dict | None = None + headers: dict = field(default_factory=dict) + reason: str | None = None + streamed: bool = False # sent already, on the stream ``open_stream`` began + + +class _EventStream: + """A request's response as Server-Sent Events, begun by the first + notification: the notifications, then the response.""" + + def __init__(self, request): + self._request = request + self._write = None + + @property + def begun(self): + return self._write is not None + + async def send(self, message): + if self._write is None: + self._write = await self._request.open_stream(EVENT_STREAM_HEADERS) + data = json.dumps(message, separators=(",", ":")) + await self._write(f"event: message\ndata: {data}\n\n") + + +def _event_stream(request): + """The event stream that a request's first notification starts, or None when + its host or its client cannot use one.""" + accept = request.headers.get("accept") or "" + if request.open_stream is None or "text/event-stream" not in accept: + return None + return _EventStream(request) + + +def _error_response(status, message, headers=None, code=SERVER_ERROR, rpc_id=None): + body = {"jsonrpc": "2.0", "id": rpc_id, "error": {"code": code, "message": message}} + return HTTPResponse(status, body, headers or {}) + + +def _decoded(value): + """A header value, decoded when sent in the Base64 form.""" + if value.startswith(BASE64_PREFIX) and value.endswith(BASE64_SUFFIX): + encoded = value[len(BASE64_PREFIX) : -len(BASE64_SUFFIX)] + try: + return base64.b64decode(encoded, validate=True).decode("utf-8") + except ValueError: # binascii.Error, UnicodeDecodeError + return None + return value + + +def _header_mismatch(headers, message): + """What is wrong with the headers a 2026-07-28 request mirrors from its + body, or None.""" + method = message.get("method") + params = message.get("params") + params = params if isinstance(params, dict) else {} + expected = [ + ("MCP-Protocol-Version", request_version(message)), + ("Mcp-Method", method), + ] + if method in NAMED_BY: + expected.append(("Mcp-Name", params.get(NAMED_BY[method]))) + for name, value in expected: + sent = headers.get(name.lower()) + if sent is None: + return f"Header mismatch: the {name} header is missing" + if name == "Mcp-Name": + sent = _decoded(sent) + if sent != value: + return ( + f"Header mismatch: {name} header value {sent!r} does not match " + f"body value {value!r}" + ) + return None + + +async def _cancel_on_disconnect(request, task): + """Cancel ``task`` if the client of ``request`` disconnects before it ends, and + return whether it did. aiohttp keeps a handler running after its client is + gone, and exposes no callback for it once the body is read.""" + while not task.done(): + await asyncio.sleep(DISCONNECT_POLL) + if request.disconnected(): + task.cancel() + return True + return False + + +class MCPEndpoint: + """Serves MCP over HTTP for a trame server (its root), whether or not + ``server.mcp`` exists yet.""" + + def __init__(self, server): + self._server = server.root_server + + @property + def _mcp(self): + return self._server._mcp + + def _origin_allowed(self, origin, request): + mcp = self._mcp + if mcp is not None and origin in mcp.http.allowed_origins: + return True + try: + parsed = urlsplit(origin) + except ValueError: + return False + if not parsed.hostname: # e.g. "null" from sandboxed frames / file:// + return False + return is_loopback(parsed.hostname) or parsed.netloc == request.authority + + def _check(self, request): + origin = request.headers.get("origin") + if origin is not None and not self._origin_allowed(origin, request): + logger.warning("Rejected an MCP request from origin %r", origin) + return _error_response(403, "Origin not allowed") + + mcp = self._mcp + token = mcp.http.token if mcp is not None else None + if token: + scheme, _, credentials = request.headers.get("authorization", "").partition( + " " + ) + if scheme.lower() != "bearer" or not hmac.compare_digest( + credentials.strip().encode(), token.encode() + ): + return _error_response( + 401, "Unauthorized", headers={"WWW-Authenticate": "Bearer"} + ) + elif not is_loopback(request.host): + return _error_response( + 403, + "Host not allowed: set a token (mcp.http.token, --mcp-http-token or " + "TRAME_MCP_HTTP_TOKEN) to serve MCP under a non-loopback host name", + ) + + if mcp is None or not mcp.enabled: + return _error_response( + 503, "MCP access is currently disabled by the application" + ) + if "http" not in mcp.transports: + return _error_response(503, "The application turned its HTTP transport off") + + if request.content_type != "application/json": + return _error_response(415, "Content-Type must be application/json") + + return None + + def _check_modern(self, request, message): + """The 400 or 404 response that a request from a 2026-07-28 client gets, or + None. Such a request names its version in its body, or its headers say it + should.""" + if ( + request_version(message) is None + and request.headers.get("mcp-protocol-version") not in MODERN_VERSIONS + ): + return None + rpc_id = message.get("id") + mismatch = _header_mismatch(request.headers, message) + if mismatch is not None: + return _error_response(400, mismatch, code=HEADER_MISMATCH, rpc_id=rpc_id) + error = check_request(message) + if error is not None: + return HTTPResponse(400, error) + if not self._mcp.protocol.handles(message): + return _error_response( + 404, + f"Method not found: {message.get('method')}", + code=METHOD_NOT_FOUND, + rpc_id=rpc_id, + ) + return None + + async def post(self, request: HTTPRequest) -> HTTPResponse: + rejection = self._check(request) + if rejection is not None: + return rejection + + try: + message = json.loads(await request.read()) + except (json.JSONDecodeError, UnicodeDecodeError): + return _error_response(400, "Parse error", code=PARSE_ERROR) + + if isinstance(message, list): + return _error_response( + 400, "JSON-RPC batching is not supported", code=INVALID_REQUEST + ) + + protocol = self._mcp.protocol + is_request = ( + isinstance(message, dict) and "method" in message and "id" in message + ) + session = None + if is_request and (rejection := self._check_modern(request, message)): + return rejection + if not isinstance(message, dict) or request_version(message) is None: + # Before 2026-07-28 only: a session scoping cancellation + session = request.headers.get(SESSION_HEADER.lower()) + stream = _event_stream(request) if is_request else None + if is_request: + notify = stream.send if stream is not None else None + handling = asyncio.ensure_future( + protocol.handle(message, session=session, notify=notify) + ) + watching = asyncio.ensure_future(_cancel_on_disconnect(request, handling)) + try: + response = await handling + except asyncio.CancelledError: + if not ( + watching.done() and not watching.cancelled() and watching.result() + ): + raise + logger.info( + "MCP client disconnected, cancelled %s request %r", + message.get("method"), + message.get("id"), + ) + # Never delivered; for the access log, as nginx does + begun = stream is not None and stream.begun + return HTTPResponse(499, reason="Client Closed Request", streamed=begun) + finally: + watching.cancel() + else: + response = await protocol.handle(message, session=session) + if response is None: + # Notification or client response: acknowledge without a body + return HTTPResponse(202) + if stream is not None and stream.begun: + await stream.send(response) + return HTTPResponse(200, streamed=True) + + headers = {} + if is_request and message["method"] == "initialize" and "result" in response: + headers[SESSION_HEADER] = secrets.token_urlsafe(24) + return HTTPResponse(200, response, headers) + + @staticmethod + def unsupported(_request=None) -> HTTPResponse: + # Stateless server: no SSE stream to offer (GET), no session to end (DELETE) + return HTTPResponse(405, headers={"Allow": "POST"}) + + +# ----------------------------------------------------------------------------- +# aiohttp +# ----------------------------------------------------------------------------- + + +def _from_aiohttp(request, open_stream=None): + def disconnected(): + transport = request.transport + return transport is None or transport.is_closing() + + return HTTPRequest( + method=request.method, + headers={k.lower(): v for k, v in request.headers.items()}, + host=request.url.host, + authority=request.host, + read=request.read, + disconnected=disconnected, + open_stream=open_stream, + ) + + +def _to_aiohttp(response): + from aiohttp import web # noqa: PLC0415 + + if response.body is None: + return web.Response( + status=response.status, headers=response.headers, reason=response.reason + ) + return web.json_response( + response.body, status=response.status, headers=response.headers + ) + + +def aiohttp_routes(endpoint, path): + """The aiohttp routes serving ``endpoint`` at ``path``.""" + from aiohttp import web # noqa: PLC0415 + + async def post(request): + streams = [] + + async def open_stream(headers): + stream = web.StreamResponse(status=200, headers=headers) + await stream.prepare(request) + streams.append(stream) + + async def write(text): + await stream.write(text.encode()) + + return write + + response = await endpoint.post(_from_aiohttp(request, open_stream)) + if not response.streamed: + return _to_aiohttp(response) + (stream,) = streams + with contextlib.suppress(ConnectionError): # the client may be gone + await stream.write_eof() + return stream + + async def other(_request): + return _to_aiohttp(endpoint.unsupported()) + + return [web.post(path, post), web.get(path, other), web.delete(path, other)] + + +def mount(server, wslink_server): + """Mount the endpoint of ``server`` on the web server wslink started, when it is + aiohttp's, and return whether it did. Other hosts can serve + :class:`MCPEndpoint` themselves.""" + from aiohttp import web # noqa: PLC0415 (wslink depends on aiohttp) + + app = getattr(wslink_server, "app", None) + mcp = server._mcp + if not isinstance(app, web.Application): + # Only worth mentioning to an app that uses agents + log = logger.info if mcp is not None else logger.debug + log( + "The MCP endpoint is not mounted: this web server is not aiohttp; " + "serve trame_server.mcp.endpoint.MCPEndpoint yourself" + ) + return False + + path = mcp.http.path if mcp is not None else "/mcp" + app.add_routes(aiohttp_routes(MCPEndpoint(server), path)) + return True + + +def generate_token(): + return secrets.token_urlsafe(32) diff --git a/src/trame_server/mcp/objects.py b/src/trame_server/mcp/objects.py new file mode 100644 index 0000000..bdf7c01 --- /dev/null +++ b/src/trame_server/mcp/objects.py @@ -0,0 +1,1815 @@ +"""The objects agents reach, and the six tools that work on them. + +Objects are: + +- instances of ``MCPExposable`` classes, such as components + (``TrameComponent``) and trame-dataclass models, whose class declares + something. Agents reach them by the names the application + gives them, or by id through the visible properties of an object they + already reach; +- the ``state`` object, whose properties are the state keys the + application exposes (``state.mcp_expose()``). + +``object_list``, ``object_types``, ``object_get``, ``object_update``, +``object_call`` and ``object_batch`` serve all objects, so the tool list +never changes as objects come and go. +""" + +from __future__ import annotations + +import asyncio +import inspect +import json +import logging +import re +import weakref +from dataclasses import dataclass, field +from typing import Annotated, Literal + +from trame_server.mcp import budget, progress, typeinfo +from trame_server.mcp.permissions import PermissionDeniedError +from trame_server.mcp.schema import ( + ToolError, + ValidationError, + bind_arguments, + error_text, + from_json, + function_input_schema, + is_object, + object_classes, + resolve_object_arguments, + return_hint, + type_to_schema, + validate, +) +from trame_server.mcp.typeinfo import ( + STATE_TYPE, + Property, + TypeInfo, + is_visible, + property_schema, + reference_schema, + subclasses, + type_info, + type_name, +) + +logger = logging.getLogger("trame_server.mcp") + +__all__ = ["OBJECT_TOOLS", "BatchError", "NotFoundError", "ObjectLayer", "StateObject"] + +OBJECT_TOOLS = ( + "object_list", + "object_types", + "object_get", + "object_update", + "object_call", + "object_batch", +) +STATE_NAME = "state" +NAME = re.compile(r"^[A-Za-z_]") +MAX_DEPTH = 10 +# Objects the connection instructions name +MAX_LISTED = 20 +# The object tools' parameters, as their input schemas say +Target = Annotated[str, "Object name or id"] +Depth = Annotated[int, {"minimum": 1, "maximum": MAX_DEPTH}] +Offset = Annotated[int, {"minimum": 0}] +TypeNames = Annotated[ + list[str] | None, + 'Types ("Viewer") or members ("Viewer.reset")', +] + +LIST_DESCRIPTION = """\ +List the objects this app names, as {"objects": {name: {"type", "read_only"?, \ +"description"?}}}. "description" says what the object is for; "type" is null \ +while the name holds nothing. The objects they refer to are reached by id, as \ +object_get shows them. When "more" is present, it gives the call that lists \ +the next ones.""" + +TYPES_DESCRIPTION = """\ +Describe the types of this app's objects. Without `type_names`, it returns an \ +index of the types; with them, types in full ("Viewer") or single members \ +("Viewer.set_view_up"). A subclass lists what it adds, under "extends", unless \ +`include_inherited`; false marks what it lacks. \ +{"object": "T"} is a value holding an object of type T: read as a reference, \ +returned by a method as object_get shows it, written by its id.""" + +GET_DESCRIPTION = """\ +Read an object. `target` is an object name or id. \ +The reply is {"result", "objects"}: "result" refers to what was read, and \ +"objects" holds each object read, by id, once, as {"type", "values"}. A \ +reference is {"id"}, and {"id", "type"} for an object not read here. `depth` \ +reads objects up to that many references away. `properties` reads only those, \ +in full. What is cut short is listed under "more", with the call that reads \ +the rest.""" + +UPDATE_DESCRIPTION = """\ +Write properties of an object: `values` maps property names to values \ +(objects by id), as object_get reads them. Every value is checked before any \ +is applied.""" + +CALL_DESCRIPTION = """\ +Call a method of an object. `args` maps parameter names to values (objects by \ +id or name). A result holding objects comes back as object_get shows it.""" + +BATCH_DESCRIPTION = """\ +Run steps in order, in one call: {"new": "", "args", "as"} builds an \ +object; {"call": target, "method", "args", "as"}; {"update": target, "values"}. \ +"@name" is the result of the step named so ("@name.prop.0" follows properties \ +and indexes, "@@" is a literal "@"). Steps are checked before any runs; the \ +first failure stops the batch and earlier steps stay applied. An object built \ +and never attached is dropped, listed in "unattached". Objects are read as \ +object_get does; with `results` "links", only their references to others.""" + +_STEP_KEYS = { + "new": ("new", "args", "as"), + "call": ("call", "method", "args", "as"), + "update": ("update", "values"), +} +_STEP_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + +_HIDDEN = object() + + +async def _settle(obj): + """Wait for the work a change started, such as the watchers trame-dataclass runs + in a task. Without this, the reply and the agent's next read would not reflect + what that work derives.""" + if not is_object(obj): + return + try: + await obj._mcp_settle() + except Exception: + # The change itself is applied: a failing watcher is the app's bug, + # not the agent's, and must not read as "nothing was applied". + logger.exception("A watcher of %s %s raised", type(obj).__name__, obj._id) + + +# ----------------------------------------------------------------------------- +# The state object +# ----------------------------------------------------------------------------- + + +class StateObject: + """The ``state`` object: its properties are the trame state keys the app + exposed (``server.state.mcp_expose()``), with their real names (a child + server's prefix included).""" + + _id = STATE_NAME + + def __init__(self): + self._keys = {} # key -> Property + self._info = None + + def __bool__(self): + return bool(self._keys) + + def expose(self, key, *, type=None, schema=None, description=None, read_only=False): + if not isinstance(read_only, bool): + msg = "read_only must be True or False" + raise TypeError(msg) + var_schema, description = property_schema( + type, {"schema": schema, "description": description} + ) + self._keys[key] = Property( + name=key, + origin="state", + schema=var_schema, + description=description, + writable=not read_only, + type_def=type, + ) + self._info = None + + def unexpose(self, *keys): + for key in keys: + self._keys.pop(key, None) + self._info = None + + @property + def info(self) -> TypeInfo: + if self._info is None: + self._info = TypeInfo(cls=None, visible=True, properties=dict(self._keys)) + return self._info + + +# ----------------------------------------------------------------------------- +# Exposure +# ----------------------------------------------------------------------------- + + +class _Entry: + """What a name exposes: an object, or a getter that returns one or None and is + called again on every request. The object is held weakly when its class asks + for it, as components do.""" + + __slots__ = ("_value", "description", "read_only", "source_type") + + def __init__(self, source_type, value, description, read_only): + self.source_type = source_type + self.description = description + self.read_only = bool(read_only) + if source_type == "weak": + try: + value = weakref.ref(value) + except TypeError: # no __weakref__: held + source_type = self.source_type = "object" + self._value = value + + def resolve(self): + if self.source_type in ("weak", "getter"): + return self._value() + return self._value + + def holds(self, obj): + """Whether this entry names ``obj`` itself.""" + return self.source_type in ("weak", "object") and self.resolve() is obj + + def exposes(self, source): + """Whether ``source`` is what this entry was given.""" + if self.source_type == "getter": + return self._value == source # bound methods are rebuilt on access + return self.resolve() is source + + def gone(self): + """Whether the object this entry held weakly was garbage collected.""" + return self.source_type == "weak" and self._value() is None + + +@dataclass +class Graph: + """What agents reach now. + + ``objects`` maps ids to ``(object, writable)``; ``entries`` maps names + to ``(value, read_only)``; ``names`` maps ids to the names exposing them. + """ + + objects: dict + entries: dict + names: dict = field(default_factory=dict) + + +@dataclass +class _Read: + """An object read for a reply, with its raw property values.""" + + obj: object + writable: bool + distance: int # references away from what was read + raw: dict + errors: dict + + +class _Scope: + """What one ``object_batch`` holds: the results named with ``as``, and + the models its ``new`` steps built (held until it ends).""" + + def __init__(self): + self.names = {} + self.built = {} # _id -> model + + +class NotFoundError(ToolError, LookupError): + """An object, type, member or property the agent named does not exist. The + message lists what does exist.""" + + +class BatchError(ToolError): + """An ``object_batch`` step failed. The message says which one, why, and what the + steps before it returned.""" + + +class ObjectLayer: + """The objects of one MCP server and the object tools.""" + + def __init__(self, mcp): + self._mcp = mcp + self._entries = {} # name -> _Entry + self.state = StateObject() + self._registered = False + self._warned = weakref.WeakSet() # classes warned about + + # ------------------------------------------------------------------------- + # Exposure (server.mcp) + # ------------------------------------------------------------------------- + + def remove(self, obj): + """Withdraw every name exposing ``obj``.""" + for name in [n for n, e in self._entries.items() if e.holds(obj)]: + del self._entries[name] + + def expose(self, source, name, description=None, read_only=False): + if not isinstance(name, str) or not NAME.match(name): + msg = f"An object name starts with a letter or an underscore, got {name!r}" + raise ValueError(msg) + if name == STATE_NAME: + msg = ( + f"{STATE_NAME!r} is reserved for the state object " + f"(server.state.mcp_expose())" + ) + raise ValueError(msg) + if description is not None and not isinstance(description, str): + msg = f"description must be a string, got {type(description).__name__}" + raise TypeError(msg) + + if is_object(source): + if not is_visible(type(source)): + cls = type(source).__name__ + msg = ( + f"{cls} declares nothing for agents: mark its members with " + f"@mcp.method / @mcp.property, or use the mcp= class keyword " + f"(or {cls}.mcp_describe() for a class you don't own)" + ) + raise TypeError(msg) + source_type = "weak" if type(source)._mcp_hold_weakly else "object" + elif callable(source) and not isinstance(source, type): + source_type = "getter" + else: + msg = ( + "expose_object() takes a component, a model, or a function " + f"returning one; got {type(source).__name__}. Expose the items " + f"of a list each under its own name, or through a property of " + f"an object." + ) + raise TypeError(msg) + + current = self._entries.get(name) + if current is not None and not current.gone() and not current.exposes(source): + if current.source_type == "getter": + held = "a function" + else: + held = self._label(current.resolve()) + msg = ( + f"{name!r} already exposes {held}; withdraw it first " + f"(server.mcp.unexpose_object({name!r}))" + ) + raise ValueError(msg) + self._entries[name] = _Entry(source_type, source, description, read_only) + self._ensure_tools() + + def unexpose(self, name): + if self._entries.pop(name, None) is None: + msg = f"No object is exposed as {name!r}" + raise KeyError(msg) + + def expose_state(self, key, **options): + self.state.expose(key, **options) + self._ensure_tools() + + def exposes_anything(self): + return bool(self._entries or self.state) + + # ------------------------------------------------------------------------- + # Reachability + # ------------------------------------------------------------------------- + + def _is_object(self, value): + return is_object(value) and is_visible(type(value)) + + def _info(self, obj) -> TypeInfo: + if obj is self.state: + return self.state.info + return type_info(type(obj)) + + def _type(self, obj) -> str: + if obj is self.state: + return STATE_TYPE + return type_name(type(obj)) + + def _label(self, obj) -> str: + if obj is self.state: + return STATE_NAME + return f"{self._type(obj)} {obj._id}" + + def _read(self, obj, prop): + if prop.origin == "state": + return self._mcp.server.state[prop.name] + return getattr(obj, prop.name) + + def _readable(self, info): + """``{name: Property}`` of the properties agents may read now: those + whose permissions are on.""" + permissions = self._mcp.permissions + return { + name: prop + for name, prop in info.properties.items() + if permissions.allows(prop.requires) + } + + def _children(self, obj): + for prop in self._readable(self._info(obj)).values(): + if not prop.reference: + continue + try: + value = self._read(obj, prop) + except Exception: + logger.debug("Reading %s.%s raised", self._type(obj), prop.name) + continue + yield from _objects_within(value, self._is_object) + + def _graph(self, scope=None) -> Graph: + """Everything agents reach now: what names expose, and the objects + visible properties refer to from there. An object is writable when a + writable path reaches it.""" + entries = {} + for name, entry in list(self._entries.items()): + try: + value = entry.resolve() + except Exception: + logger.exception("Resolving the object exposed as %r raised", name) + value = None + if value is not None and not self._is_object(value): + logger.warning( + "%r exposes a %s, not a component or model agents can see", + name, + type(value).__name__, + ) + value = None + entries[name] = (value, entry.read_only) + + names = {} # id -> the names exposing it + roots = [] + for name, (value, read_only) in entries.items(): + if value is not None: + names.setdefault(value._id, []).append(name) + roots.append((value, not read_only)) + if scope is not None: + roots.extend((obj, True) for obj in scope.built.values()) + + objects = {} + for writable in (True, False): + stack = [obj for obj, w in roots if w is writable] + while stack: + obj = stack.pop() + if obj._id in objects: + continue + objects[obj._id] = (obj, writable) + stack.extend(self._children(obj)) + if self.state: + objects[STATE_NAME] = (self.state, True) + + return Graph(objects=objects, entries=entries, names=names) + + def _unknown(self, target, graph): + names = [*graph.entries, *([STATE_NAME] if self.state else [])] + if not names: + return ( + f"Unknown object {json.dumps(target)}: no object has that name or " + f"id; object_list() lists the objects." + ) + return ( + f"Unknown object {json.dumps(target)}. Objects: {', '.join(names)} " + f"(or an object id); object_list() lists them all." + ) + + def _object(self, target, graph): + """``(object, writable)`` for a target name or id.""" + if not isinstance(target, str) or not target: + msg = "`target` must be an object name or id" + raise ValidationError(msg) + if target in graph.entries: + value, _ = graph.entries[target] + if value is None: + msg = f"The object exposed as {json.dumps(target)} is currently empty" + raise NotFoundError(msg) + found = graph.objects.get(value._id) + if found is not None: + return found + found = graph.objects.get(target) + if found is not None: + return found + raise NotFoundError(self._unknown(target, graph)) + + def _resolver(self, scope, read_only): + """``resolve(value, cls, path)`` for object-typed parameters and + properties: an id or name of a reachable object of class ``cls``.""" + graph = None + + def resolve(value, cls, path): + nonlocal graph + if graph is None: + graph = self._graph(scope) + found = graph.objects.get(value) + if found is None and value in graph.entries: + named = graph.entries[value][0] + if named is not None: + found = graph.objects.get(named._id) + expected = type_name(cls) + if found is None or found[0] is self.state: + msg = ( + f"{path}: {json.dumps(value)} is not the id or name of an object " + f"agents can reach (an object of type {expected} is expected)" + ) + raise ValidationError(msg) + obj, writable = found + if not isinstance(obj, cls): + msg = ( + f"{path}: {json.dumps(value)} has type {self._type(obj)}, " + f"not {expected}" + ) + raise ValidationError(msg) + if not writable and not read_only: + msg = ( + f"{path}: {self._label(obj)} is exposed read-only; only read-only " + f"methods can take it" + ) + raise PermissionError(msg) + return obj + + return resolve + + def resolve_arguments( + self, fn, arguments, *, read_only=False, scope=None, label="arguments" + ): + """``arguments`` of ``fn``, with the ids or names given for its object-typed + parameters replaced by the objects. An object must be reachable (or built + by the batch ``scope``) and of the annotated class, and one reached only + read-only can be passed to ``read_only`` callables only.""" + return resolve_object_arguments( + fn, arguments, self._resolver(scope, read_only), label + ) + + # ------------------------------------------------------------------------- + # Replies: what was read, with its objects as references, and the + # objects read, once each, in a flat table + # ------------------------------------------------------------------------- + + def _reply( + self, + value, + graph, + depth=1, + *, + target=None, + writable=True, + properties=None, + links=False, + spread=False, + ): + """``{"result", "objects"?, "more"?}``: ``value`` with its objects as + references, and the objects within ``depth`` steps of it read into + ``objects``, by id, nearest first. A reference is ``{"id"}`` for an object + that is read, and ``{"id", "type"}`` for one that is not. + + Over the budget, the farthest objects are left out first, since + object_get reads them again; ``value`` is cut last. + + :param target: the name the agent used for ``value``, used in the calls + that notes suggest. + :param writable: whether ``value``, a target object, is writable. + :param properties: the only properties of ``value`` to read, which + come back whole. + :param links: read only the properties holding objects. + :param spread: ``value``, a dict, gives the reply's own fields in + place of ``result``. + """ + reads = self._gather(value, graph, depth, writable, properties) + order = list(reads) + requested = set(properties or ()) + + def call_target(obj_id): + return target if target and reads[obj_id].obj is value else obj_id + + def build(count): + read = set(order[:count]) + more = {} + objects = {} + for obj_id in order[:count]: + entry, notes = self._entry( + reads[obj_id], + graph, + reads, + read, + (call_target(obj_id), depth), + links, + requested, + ) + objects[obj_id] = entry + more.update(notes) + if count < len(order): + more["objects"] = ( + f"{count} of {len(order)} objects read to fit the response " + f"budget; object_get(id) reads one" + ) + result = self._encode(value, graph, reads, read) + reply = {"result": None if result is _HIDDEN else result} + if objects: + reply["objects"] = objects + if more: + reply["more"] = more + return reply + + def finish(reply): + if spread: + return {**reply.pop("result"), **reply} + return reply + + room = self._budget() + reply = build(len(order)) + if budget.size(reply) <= room: + return finish(reply) + # Leave out the farthest objects: the most that fit. One object read + # for itself stays, and its values are cut instead. + single = bool(order) and reads[order[0]].obj is value + low, high = int(single), len(order) - 1 + while low < high: + middle = (low + high + 1) // 2 + if budget.size(build(middle)) <= room: + low = middle + else: + high = middle - 1 + reply = build(low) + if budget.size(reply) <= room: + return finish(reply) + if single: # the object read alone is still too big + return self._fit_entry( + reply, reads[order[0]], requested, (call_target(order[0]), depth) + ) + # Then the result, noted under the fields the agent sees + result = reply["result"] if spread else {"result": reply["result"]} + rest = budget.size({k: v for k, v in reply.items() if k != "result"}) + result, cuts = budget.shrink(result, room - rest - 200) + reply["result"] = result if spread else result["result"] + more = reply.setdefault("more", {}) + for key, note in cuts.items(): + more[key] = f"{note} to fit the response budget" + return finish(reply) + + def _gather(self, value, graph, depth, writable, properties): + """``{id: _Read}``: the objects ``value`` holds, when agents can see their type, + then those reachable within ``depth`` steps, nearest first.""" + reads = {} + pending = [] + + def admit(obj, distance): + if obj._id in reads or distance >= depth: + return + found = graph.objects.get(obj._id) + if distance == 0: # what was read: the target, or a call's result + if obj is not self.state and not is_visible(type(obj)): + return + if obj is value: + obj_writable = writable + else: + obj_writable = found[1] if found and found[0] is obj else True + elif found is None or found[0] is not obj: + return + else: + obj_writable = found[1] + wanted = properties if obj is value else None + reads[obj._id] = self._read_object(obj, obj_writable, distance, wanted) + pending.append(obj._id) + + def visit(raw, distance, capped): + # A capped value keeps only its first items (see _entry). The + # objects in the other items are not shown, so they are not read + limit = budget.MAX_ITEMS if capped else None + for obj in _objects_within(raw, obj_like, limit): + admit(obj, distance) + + def obj_like(raw): + return raw is self.state or is_object(raw) + + visit(value, 0, False) + requested = set(properties or ()) + index = 0 + while index < len(pending): + read = reads[pending[index]] + index += 1 + for name, raw in read.raw.items(): + whole = read.distance == 0 and name in requested + visit(raw, read.distance + 1, not whole) + return reads + + def _read_object(self, obj, writable, distance, wanted=None): + raw, errors = {}, {} + for prop in self._readable(self._info(obj)).values(): + if wanted is not None and prop.name not in wanted: + continue + try: + raw[prop.name] = self._read(obj, prop) + except Exception as e: + errors[prop.name] = error_text(e) + return _Read(obj, writable, distance, raw, errors) + + def _encode(self, value, graph, reads, read): + """``value`` as JSON: objects as references, those agents cannot + reach left out (``_HIDDEN``).""" + if value is self.state or is_object(value): + if value._id in read: + return {"id": value._id} + found = graph.objects.get(value._id) + known = reads.get(value._id) + if (known is not None and known.obj is value) or ( + found is not None and found[0] is value + ): + return {"id": value._id, "type": self._type(value)} + return _HIDDEN + if isinstance(value, (list, tuple, set, frozenset)): + items = [self._encode(v, graph, reads, read) for v in value] + items = [v for v in items if v is not _HIDDEN] + if isinstance(value, (set, frozenset)): + return sorted(items, key=str) + return items + if isinstance(value, dict): + items = { + str(k): self._encode(v, graph, reads, read) for k, v in value.items() + } + return {k: v for k, v in items.items() if v is not _HIDDEN} + if value is None or isinstance(value, (str, int, float, bool)): + return value + tolist = getattr(value, "tolist", None) # numpy and friends + if callable(tolist): + try: + return self._encode(tolist(), graph, reads, read) + except Exception: + pass + return str(value) + + def _entry(self, found, graph, reads, read, call, links, requested): + """``(entry, more)``: an object read, as ``objects`` lists it, with its + values capped unless they were requested. ``call`` is ``(target, + depth)``, used in the calls that notes suggest.""" + obj = found.obj + target, depth = call + values, more = {}, {} + for name, raw in found.raw.items(): + rendered = self._encode(raw, graph, reads, read) + rendered = None if rendered is _HIDDEN else rendered + if links and not _holds_references(rendered): + continue + if found.distance or name not in requested: + rendered, note = budget.cap(rendered) + if note: + rest = self._get_call(target, name, depth - found.distance) + more[f"{obj._id}.{name}"] = f"{note}; {rest} for all" + values[name] = rendered + entry = {"type": self._type(obj)} + if graph.names.get(obj._id): + entry["names"] = graph.names[obj._id] + if values: + entry["values"] = values + if found.errors: + entry["errors"] = found.errors + if not found.writable: + entry["read_only"] = True + return entry, more + + def encode_tool_result(self, value): + """An app tool's result, the way object_get shows it when it holds objects, and + as is otherwise.""" + if not _holds_objects(value): + return value + return self._reply(value, self._graph()) + + def _encode_result(self, value, graph): + """A method's or constructor's result, the way object_get shows it when it holds + objects. Anything else is returned as is, cut to fit the budget.""" + if _holds_objects(value): + reply = self._reply(value, graph) + if reply["result"] is None and not reply.get("objects"): + return None + return reply + return self._fit_result(value) + + # ------------------------------------------------------------------------- + # Response budget + # ------------------------------------------------------------------------- + + def _budget(self): + return self._mcp.max_response_tokens * budget.CHARS_PER_TOKEN + + @staticmethod + def _get_call(target, name, depth=1): + call = f"object_get({json.dumps(target)}, properties=[{json.dumps(name)}]" + if depth > 1: + call += f", depth={depth}" + return f"{call})" + + def _fit_entry(self, reply, found, requested, call): + """Cut the largest values of a reply that reads one object, until it fits. Notes + go in ``reply["more"]``, keyed ``id.property``, and their counts come from the + values ``found`` read, before they were capped.""" + target, depth = call + obj_id = found.obj._id + entry = reply["objects"][obj_id] + values, more = entry.get("values", {}), reply.setdefault("more", {}) + room = self._budget() - budget.size({**reply, "objects": {}}) + room -= budget.size({obj_id: {**entry, "values": {}}}) + if budget.size(values) > room: + totals = { + name: len(raw) + for name, raw in found.raw.items() + if isinstance(raw, (str, list, tuple, set, frozenset)) + } + values, cuts = budget.shrink(values, room, totals) + for name, note in cuts.items(): + if name in requested: + more[f"{obj_id}.{name}"] = f"{note} to fit the response budget" + else: + rest = self._get_call(target, name, depth) + more[f"{obj_id}.{name}"] = f"{note}; {rest} alone for more" + while values and budget.size(values) > room - budget.size(more): + name = max(values, key=lambda k: budget.size(values[k])) + del values[name] + rest = self._get_call(target, name, depth) + more[f"{obj_id}.{name}"] = ( + f"left out to fit the response budget; {rest}" + ) + if values: + entry["values"] = values + else: + entry.pop("values", None) + return reply + + def _fit_result(self, value): + """A plain value (a method's result) cut to fit the budget.""" + room = self._budget() + if budget.size(value) <= room: + return value + wrapped, cuts = budget.shrink({"result": value}, room - 120) + note = cuts.get("result", "shortened") + return { + "result": wrapped["result"], + "more": f"{note} to fit the response budget", + } + + # ------------------------------------------------------------------------- + # object_list and object_types + # ------------------------------------------------------------------------- + + def _catalog(self, graph): + """``{type name: TypeInfo}``: the types of the objects, those their + properties may hold (with visible subclasses), and the visible bases + they extend.""" + pending = [ + type(obj) for obj, _ in graph.objects.values() if obj is not self.state + ] + known, expanded = set(), set() + while pending: + cls = pending.pop() + if cls in known or not is_visible(cls): + continue + known.add(cls) + for prop in type_info(cls).properties.values(): + if prop.reference: + for held in object_classes(prop.type_def): + if held not in expanded: + expanded.add(held) + pending.extend(subclasses(held)) + for cls in list(known): + base = self._base(cls) + while base is not None and base not in known: + known.add(base) + base = self._base(base) + catalog = {type_name(cls): type_info(cls) for cls in known} + if self.state: + catalog[STATE_TYPE] = self.state.info + return dict(sorted(catalog.items())) + + @staticmethod + def _base(klass): + """The nearest visible base of ``klass``, which it can extend.""" + for base in klass.__mro__[1:]: + if is_visible(base): + return base + return None + + def _render(self, info) -> dict: + """A type in full: what agents may use of it now.""" + permissions = self._mcp.permissions + entry = {} + if info.description: + entry["description"] = info.description + properties = {} + for prop in self._readable(info).values(): + item = prop.described() + writable = prop.writable and permissions.allows(prop.set_requires) + if not writable: + item["read_only"] = True + requires = prop.requires + (prop.set_requires if writable else ()) + if requires: + item["requires"] = list(dict.fromkeys(requires)) + properties[prop.name] = item + methods = {} + for method in info.methods.values(): + if not permissions.allows(method.requires): + continue + fn = getattr(info.cls, method.attr) + item = {"parameters": function_input_schema(fn, ref=reference_schema)} + if method.description: + item = {"description": method.description, **item} + returns = type_to_schema(return_hint(fn), reference_schema) + if returns and returns != {"type": "null"}: + item["returns"] = returns + if method.read_only: + item["read_only"] = True + if method.annotations: + item["annotations"] = dict(method.annotations) + if method.requires: + item["requires"] = list(method.requires) + methods[method.name] = item + entry["properties"] = properties + entry["methods"] = methods + constructor = info.constructor + if constructor is not None and permissions.allows(constructor.requires): + item = {"parameters": self._constructor_schema(info)} + if constructor.description: + item["description"] = constructor.description + if constructor.requires: + item["requires"] = list(constructor.requires) + entry["constructor"] = item + if info.constructor_warning and info.cls not in self._warned: + self._warned.add(info.cls) + logger.warning( + "%s has no constructor for agents: %s", + info.cls.__name__, + info.constructor_warning, + ) + return entry + + def _rendered(self, catalog, include_inherited): + """``{name: entry}``: each type in full, or as what it adds to the + type it extends.""" + full = {name: self._render(info) for name, info in catalog.items()} + result = {} + for name, info in catalog.items(): + entry = full[name] + base = ( + None if include_inherited or info.cls is None else self._base(info.cls) + ) + if base is not None: + base_name = type_name(base) + if base_name not in full: + full[base_name] = self._render(type_info(base)) + entry = _extending(entry, base_name, full[base_name]) + result[name] = entry + return result + + def _named(self, graph, short=False): + """``{name: {"type", "read_only"?, "description"?}}`` for the names agents + start from, in the order they were given. ``type`` is None while a name + holds nothing, and with ``short`` the description is only its first + sentence.""" + named = {} + for name, (value, read_only) in graph.entries.items(): + item = {"type": None if value is None else self._type(value)} + if read_only: + item["read_only"] = True + about = self._entries[name].description + if about and short: + about = _first_sentence(about) + if about: + item["description"] = about + named[name] = item + if self.state: + named[STATE_NAME] = {"type": STATE_TYPE} + return named + + def list_objects(self, offset: Offset = 0): + items = list(self._named(self._graph()).items()) + start = offset + # As many as fit, one at least, with room for the note + room = self._budget() - 200 + used, end = budget.size({"objects": {}}), start + while end < len(items): + name, item = items[end] + used += budget.size({name: item}) + if used > room and end > start: + break + end += 1 + result = {"objects": dict(items[start:end])} + if end < len(items): + result["more"] = { + "objects": ( + f"{start + 1}-{end} of {len(items)} listed to fit the response " + f"budget; object_list(offset={end}) lists the next ones" + ) + } + return result + + def types(self, type_names: TypeNames = None, include_inherited: bool = False): + graph = self._graph() + catalog = self._catalog(graph) + if type_names is None: + rendered = self._rendered(catalog, include_inherited) + index = {name: _index_entry(e) for name, e in rendered.items()} + return self._fit_index({"types": index}) + + known = {**typeinfo.type_names()[1], **dict.fromkeys(catalog)} + wanted, members, unknown = [], [], [] + for raw in dict.fromkeys(type_names): + if raw in known: + wanted.append(raw) + else: + owner, _, member = raw.rpartition(".") + if owner in known and member: + members.append((owner, member)) + else: + unknown.append(raw) + + def info_of(name): + return catalog.get(name) or type_info(typeinfo.type_names()[1][name]) + + found, missing = [], [] + for owner, member in members: + entry = self._render(info_of(owner)) + section = next( + (k for k in ("properties", "methods") if member in entry[k]), None + ) + if section is not None: + found.append((owner, section, member, entry[section][member])) + elif member == "constructor" and "constructor" in entry: + found.append((owner, None, member, entry["constructor"])) + else: + missing.append(f"{owner}.{member}") + # Every name that is wrong, in one reply + problems = [] + if unknown: + problems.append( + f"Unknown types or members: {', '.join(unknown)}. Types: " + f"{', '.join(catalog)}; object_types() lists them." + ) + if missing: + problems.append( + f"No member {', '.join(missing)}; object_types([type]) lists a " + f"type's properties and methods." + ) + if problems: + raise NotFoundError(" ".join(problems)) + + selected = {name: info_of(name) for name in dict.fromkeys(wanted)} + rendered = self._rendered(selected, include_inherited) if selected else {} + for owner, section, member, value in found: + if section is None: + rendered.setdefault(owner, {})[member] = value + else: + rendered.setdefault(owner, {}).setdefault(section, {})[member] = value + result = {"types": rendered} + if budget.size(result) > self._budget(): + types, cuts = budget.shrink(rendered, self._budget() - 200) + result = {"types": types} + if cuts: + result["more"] = dict.fromkeys( + cuts, + "shortened to fit the response budget; describe fewer types per call", + ) + return result + + def _fit_index(self, result): + room = self._budget() + if budget.size(result) <= room: + return result + for entry in result["types"].values(): + entry.pop("properties", None) + entry.pop("methods", None) + note = ( + 'member names left out to fit the response budget; object_types(["Type"]) ' + "for a type's members" + ) + if budget.size(result) > room: + types = list(result["types"].items()) + keep = len(types) + while keep > 1 and budget.size(dict(types[:keep])) > room - 200: + keep //= 2 + result["types"] = dict(types[:keep]) + note = ( + f"{keep} of {len(types)} types shown, member names left out, to fit " + f'the response budget; object_types(["Type"]) for one type' + ) + result["more"] = {"types": note} + return result + + def describe_class(self, cls): + """What agents see of ``cls`` now (``server.mcp.describe(cls)``).""" + if cls is None or cls == STATE_TYPE: + return self._render(self.state.info) + info = type_info(cls) + if not info.visible: + return None + return self._render(info) + + # ------------------------------------------------------------------------- + # object_get + # ------------------------------------------------------------------------- + + def get( + self, target: Target, properties: list[str] | None = None, depth: Depth = 1 + ): + graph = self._graph() + value, writable = self._object(target, graph) + if properties is not None: + visible = list(self._readable(self._info(value))) + unknown = [p for p in properties if p not in visible] + if unknown: + msg = ( + f"{self._label(value)} has no properties {unknown}; its properties: " + f"{', '.join(visible)}" + ) + raise NotFoundError(msg) + return self._reply( + value, + graph, + depth, + target=target, + writable=writable, + properties=properties, + ) + + # ------------------------------------------------------------------------- + # object_update + # ------------------------------------------------------------------------- + + async def update(self, target: Target, values: dict): + obj, writable = self._object(target, self._graph()) + return await self._apply_update(obj, writable, values) + + async def _apply_update(self, obj, writable, values, scope=None): + label = self._label(obj) + if not writable: + msg = f"Nothing was applied: {label} is exposed read-only." + raise PermissionError(msg) + + info, permissions = self._info(obj), self._mcp.permissions + props = self._readable(info) + rejected, read_only = [], False + for name in values: + prop = props.get(name) + if prop is None: + rejected.append(f"{name} (unknown)") + elif not prop.writable: + rejected.append(f"{name} (read-only)") + read_only = True + elif not permissions.allows(prop.set_requires): + off = [n for n in prop.set_requires if not permissions.allows((n,))] + rejected.append( + f"{name} (read-only: permission {', '.join(off)} is off)" + ) + read_only = True + if rejected: + writable_names = [ + n + for n, p in props.items() + if p.writable and permissions.allows(p.set_requires) + ] + msg = ( + f"Nothing was applied. Rejected on {label}: {', '.join(rejected)}. " + f"Writable: {', '.join(writable_names) or 'none'}" + ) + # Methods may change what can't be written: point at them + methods = [ + m.name for m in info.methods.values() if permissions.allows(m.requires) + ] + if read_only and methods: + msg += ( + f". Methods: {', '.join(methods)}; " + f"object_types([{json.dumps(self._type(obj))}]) says what " + f"they do." + ) + raise ValidationError(msg) + + converted, errors = {}, [] + resolve = self._resolver(scope, read_only=False) + for name, value in values.items(): + prop = props[name] + try: + converted[name] = from_json( + validate(value, prop.schema, name), + prop.type_def, + resolve if prop.reference else None, + name, + ) + except (ValidationError, PermissionError) as e: + errors.append(str(e)) + if errors: + msg = f"Nothing was applied: {'; '.join(errors)}" + raise ValidationError(msg) + + state = self._mcp.server.state + applied = [] + with state: + # Fields and state keys first, all at once: they cannot fail + for name, value in converted.items(): + if props[name].origin == "state": + state[name] = value + applied.append(name) + elif props[name].origin == "field": + setattr(obj, name, value) + applied.append(name) + # Then setters, in order: one that raises stops the update + for name, value in converted.items(): + if props[name].origin != "property": + continue + try: + setattr(obj, name, value) + except Exception as e: + done = f"; applied: {', '.join(applied)}" if applied else "" + msg = f"Setting {label}.{name} failed: {error_text(e)}{done}" + raise ToolError(msg) from e + applied.append(name) + await _settle(obj) + return {"id": obj._id, "updated": list(converted)} + + # ------------------------------------------------------------------------- + # object_call + # ------------------------------------------------------------------------- + + def _method(self, obj, writable, method, target): + info = self._info(obj) + permissions = self._mcp.permissions + found = info.methods.get(method) + if found is None: + type_name = self._type(obj) + msg = ( + f"{target} (type {type_name}) has no method {json.dumps(method)}; " + f'object_types(["{type_name}"]) lists its methods.' + ) + raise NotFoundError(msg) + if not writable and not found.read_only: + allowed = [n for n, m in info.methods.items() if m.read_only] + msg = ( + f"{self._label(obj)} is exposed read-only; only read-only methods can " + f"be called: {', '.join(allowed) or 'none'}" + ) + raise PermissionError(msg) + permissions.require(*found.requires, what=f"{self._type(obj)}.{found.name}") + return found + + async def call(self, target: Target, method: str, args: dict | None = None): + # object_call skips the protocol's lock: the method's own concurrent + # option decides, so it is looked up first, then again under the + # lock, where the objects are as earlier calls left them. + obj, writable = self._object(target, self._graph()) + found = self._method(obj, writable, method, target) + async with self._mcp.protocol.locked(found.concurrent): + obj, writable = self._object(target, self._graph()) + found = self._method(obj, writable, method, target) + _, encoded = await self._invoke(obj, found, args, "args") + return encoded + + async def _invoke(self, obj, method, args, label, scope=None): + """Call an exposed method: bind, resolve object ids, run, settle, and + encode what it returns.""" + fn = getattr(obj, method.attr) + schema = function_input_schema(fn, ref=reference_schema) + try: + args = bind_arguments(fn, args, schema, label=label) + except ValidationError as e: + msg = f"Invalid arguments for {self._type(obj)}.{method.name}: {e}" + raise ValidationError(msg) from None + args = self.resolve_arguments( + fn, args, read_only=method.read_only, scope=scope, label=label + ) + result = fn(**progress.inject(fn, args)) + if inspect.isawaitable(result): + result = await result + await _settle(obj) + return result, self._encode_result(result, self._graph(scope)) + + # ------------------------------------------------------------------------- + # object_batch + # ------------------------------------------------------------------------- + + def _constructor_schema(self, info): + constructor = info.constructor + if constructor.attr is not None: + fn = getattr(info.cls, constructor.attr) + return function_input_schema( + fn, skip=(constructor.server,), ref=reference_schema + ) + properties = { + prop.name: prop.described() + for prop in self._readable(info).values() + if prop.origin == "field" and prop.writable + } + return { + "type": "object", + "properties": properties, + "additionalProperties": False, + } + + def _constructible(self, graph): + """Types with a constructor, by name.""" + return { + name: info + for name, info in self._catalog(graph).items() + if info.constructor is not None + } + + @staticmethod + def _references(value): + """Every "@..." reference string inside a step value.""" + if isinstance(value, str): + if value.startswith("@") and not value.startswith("@@"): + yield value + elif isinstance(value, list): + for item in value: + yield from ObjectLayer._references(item) + elif isinstance(value, dict): + for item in value.values(): + yield from ObjectLayer._references(item) + + def _check_steps(self, steps): + """Validate every step before any runs, and return the types that can + be built.""" + + def fail(message): + msg = f"Nothing was applied: {message}" + raise ValidationError(msg) + + graph = self._graph() + constructible = None + defined = set() + built = {} # "as" name of a new step -> its TypeInfo + for index, step in enumerate(steps): + path = f"steps[{index}]" + step_types = [step_type for step_type in _STEP_KEYS if step_type in step] + if len(step_types) != 1: + fail(f"{path}: a step has exactly one of 'new', 'call' or 'update'") + step_type = step_types[0] + unexpected = sorted(set(step) - set(_STEP_KEYS[step_type])) + if unexpected: + fail( + f"{path}: unexpected {unexpected}; {step_type} steps take " + f"{list(_STEP_KEYS[step_type])}" + ) + if not isinstance(step[step_type], str) or not step[step_type]: + fail(f"{path}.{step_type}: expected a non-empty string") + if step_type == "call" and ( + not isinstance(step.get("method"), str) or not step["method"] + ): + fail(f"{path}: a call step needs 'method'") + if step_type == "update" and not isinstance(step.get("values"), dict): + fail( + f"{path}: an update step needs 'values', an object of property values" + ) + if step.get("args") is not None and not isinstance(step["args"], dict): + fail(f"{path}.args: expected an object") + if step_type == "new": + if constructible is None: + constructible = self._constructible(graph) + name = step["new"] + if name not in constructible: + fail( + f"{path}: no type {json.dumps(name)} can be built; types with a " + f"constructor: {', '.join(constructible) or 'none'}" + ) + info = constructible[name] + self._check_permissions( + info.constructor.requires, f"{path} (new {name})" + ) + if "as" in step: + built[step["as"]] = info + checked = [step.get("args")] + else: + checked = [step[step_type], step.get("args"), step.get("values")] + info = self._info_before_run(step[step_type], built, graph) + if info is not None and step_type == "call": + method = info.methods.get(step["method"]) + if method is not None: + what = f"{path} (call {step['method']} on {step[step_type]})" + self._check_permissions(method.requires, what) + if info is not None and step_type == "update": + for prop_name in step["values"]: + prop = info.properties.get(prop_name) + if prop is not None: + what = f"{path} (update {prop_name} on {step[step_type]})" + self._check_permissions( + prop.requires + prop.set_requires, what + ) + for reference in self._references(checked): + if reference[1:].split(".")[0] not in defined: + fail( + f'{path}: {reference!r} names no earlier step (give that step "as")' + ) + if "as" in step: + name = step["as"] + if not isinstance(name, str) or not _STEP_NAME.match(name): + fail( + f"{path}.as: expected a name of letters, digits and underscores" + ) + if name in defined: + fail(f"{path}.as: {name!r} already names an earlier step") + defined.add(name) + return constructible or {} + + def _check_permissions(self, requires, what): + try: + self._mcp.permissions.require(*requires, what=what) + except PermissionDeniedError as e: + msg = f"Nothing was applied: {e}" + raise PermissionDeniedError(msg) from None + + def _info_before_run(self, target, built, graph): + """The type of a step's target when it is known before the batch + runs: an object reachable now, or "@name" of a new step.""" + if target.startswith("@") and not target.startswith("@@"): + return built.get(target[1:]) + try: + obj, _ = self._object(target.removeprefix("@"), graph) + except LookupError: + return None + return self._info(obj) + + def _follow(self, reference, scope, path): + """The value an "@name.seg..." reference stands for.""" + name, *segments = reference[1:].split(".") + value = scope.names[name] + walked = f"@{name}" + for segment in segments: + if is_object(value) or value is self.state: + prop = self._readable(self._info(value)).get(segment) + if prop is None: + msg = ( + f"{path}: {walked} (type {self._type(value)}) has no " + f"property {segment!r}" + ) + raise NotFoundError(msg) + value = self._read(value, prop) + if prop.reference: # as object_get shows it: objects only + if isinstance(value, (list, tuple)): + value = [ + v for v in value if not is_object(v) or self._is_object(v) + ] + elif is_object(value) and not self._is_object(value): + value = None + elif isinstance(value, dict): + if segment not in value: + msg = f"{path}: {walked} has no key {segment!r}" + raise NotFoundError(msg) + value = value[segment] + elif isinstance(value, (list, tuple)): + try: + value = value[int(segment)] + except (ValueError, IndexError): + msg = ( + f"{path}: {walked} has {len(value)} items, no index {segment!r}" + ) + raise NotFoundError(msg) from None + else: + type_name = "null" if value is None else type(value).__name__ + msg = f"{path}: {walked} is {type_name}, which has no {segment!r}" + raise NotFoundError(msg) + walked = f"{walked}.{segment}" + return value + + def _substitute(self, value, scope, path, ids=True): + """``value`` with its references resolved. With ``ids``, objects become their + ids, since arguments travel as JSON.""" + if isinstance(value, str): + if value.startswith("@@"): + return value[1:] + if value.startswith("@"): + value = self._follow(value, scope, path) + return self._as_ids(value) if ids else value + return value + if isinstance(value, list): + return [ + self._substitute(v, scope, f"{path}[{i}]", ids) + for i, v in enumerate(value) + ] + if isinstance(value, dict): + return { + k: self._substitute(v, scope, f"{path}.{k}", ids) + for k, v in value.items() + } + return value + + def _as_ids(self, value): + if is_object(value) or value is self.state: + return value._id + if isinstance(value, (list, tuple)): + return [self._as_ids(v) for v in value] + if isinstance(value, dict): + return {k: self._as_ids(v) for k, v in value.items()} + return value + + def _target(self, target, scope, path): + """``(object, writable)`` for a step's target: a name, an id, or a + reference to an object.""" + target = self._substitute(target, scope, path, ids=False) + graph = self._graph(scope) + if is_object(target) or target is self.state: + found = graph.objects.get(target._id) + if found is None or found[0] is not target: + msg = f"{path}: {self._label(target)} is not reachable" + raise NotFoundError(msg) + return found + if not isinstance(target, str): + msg = f"{path}: the target is {type(target).__name__}, not an object" + raise NotFoundError(msg) + return self._object(target, graph) + + async def _build(self, info, args, scope, label): + cls, constructor = info.cls, info.constructor + self._mcp.permissions.require( + *constructor.requires, what=f"{type_name(cls)} constructor" + ) + schema = self._constructor_schema(info) + if constructor.attr is None: + values = validate(args or {}, schema, label) + resolve = self._resolver(scope, read_only=False) + for name, given in values.items(): + prop = info.properties[name] + values[name] = from_json( + given, + prop.type_def, + resolve if prop.reference else None, + f"{label}.{name}", + ) + instance = cls(self._mcp.server, **values) + else: + fn = getattr(cls, constructor.attr) + args = bind_arguments(fn, args, schema, label=label, positional=(None,)) + args = self.resolve_arguments(fn, args, scope=scope, label=label) + instance = fn(self._mcp.server, **args) + if inspect.isawaitable(instance): + instance = await instance + if not isinstance(instance, cls): + msg = ( + f"{cls.__name__}.{constructor.attr} returned an object of type " + f"{type(instance).__name__}, not {cls.__name__}" + ) + raise TypeError(msg) + scope.built[instance._id] = instance + await _settle(instance) + return instance + + async def _run_step(self, step_type, step, scope, path, constructible): + """Run one step and return its result, which the step's ``as`` can name.""" + if step_type == "new": + info = constructible[step["new"]] + label = f"{path}.args" + args = self._substitute(step.get("args"), scope, label) + return await self._build(info, args, scope, label) + + obj, writable = self._target(step[step_type], scope, f"{path}.{step_type}") + if step_type == "update": + values = self._substitute(step["values"], scope, f"{path}.values") + return await self._apply_update(obj, writable, values, scope) + method = self._method(obj, writable, step["method"], step[step_type]) + args = self._substitute(step.get("args"), scope, f"{path}.args") + raw, _ = await self._invoke(obj, method, args, f"{path}.args", scope) + return raw + + async def batch( + self, steps: list[dict], results: Literal["links", "full"] = "links" + ): + links = results == "links" + constructible = self._check_steps(steps) + scope = _Scope() + step_results = [] + # Progress counts steps, and a method's own reports count within its step + reporter = progress.current() + for index, step in enumerate(steps): + if index: + await asyncio.sleep(0) # a cancel lands between steps + step_type = next(k for k in _STEP_KEYS if k in step) + path = f"steps[{index}]" + try: + with progress.reporting(progress.batch_step(index, len(steps))): + raw = await self._run_step( + step_type, step, scope, path, constructible + ) + except Exception as e: + applied = None + if step_results: + applied = self._reply( + {"results": step_results}, + self._graph(scope), + links=links, + spread=True, + ) + failure = _batch_failure(index, step_type, step, e, applied) + raise BatchError(failure) from e + step_results.append(raw) + if "as" in step: + scope.names[step["as"]] = raw + if reporter is not None: + await reporter(index + 1, total=len(steps)) + graph = self._graph() + unattached = [ + instance + for instance in scope.built.values() + if instance._id not in graph.objects + ] + result = {"results": step_results} + if unattached: + result["unattached"] = unattached + # Objects as they are once every step ran, each read once + reply = self._reply(result, self._graph(scope), links=links, spread=True) + for key in ("results", "unattached"): + if key in reply.get("more", {}): + reply["more"][key] += "; run fewer steps per batch" + return reply + + # ------------------------------------------------------------------------- + # Registration, instructions, audit + # ------------------------------------------------------------------------- + + def _ensure_tools(self): + if self._registered: + return + self._registered = True + + add = self._mcp._add_builtin_tool + read = {"readOnlyHint": True} + add(self.list_objects, "object_list", LIST_DESCRIPTION, read) + add(self.types, "object_types", TYPES_DESCRIPTION, read) + add(self.get, "object_get", GET_DESCRIPTION, read) + add(self.update, "object_update", UPDATE_DESCRIPTION, {"idempotentHint": True}) + # call() takes the lock itself, so that each method's option decides + add(self.call, "object_call", CALL_DESCRIPTION, concurrent=True) + # The whole batch holds the call lock, once + add(self.batch, "object_batch", BATCH_DESCRIPTION) + + def instructions(self): + """The lines the connection instructions give about the objects.""" + try: + named = self._named(self._graph(), short=True) + except Exception: + logger.exception("Listing the objects for the instructions failed") + return [] + if not named: + return [] + + def label(item): + text = item["type"] or "empty" + if item.get("read_only"): + text += ", read-only" + if item.get("description"): + text += f": {item['description']}" + return text + + items = list(named.items()) + shown = ", ".join( + f"{name} ({label(item)})" for name, item in items[:MAX_LISTED] + ) + if len(items) > MAX_LISTED: + shown += f", and {len(items) - MAX_LISTED} more" + return [ + f"Objects: {shown}; object_list() lists them, object_types() describes " + f"their types." + ] + + def start(self): + """Called when the server starts: settle the type names, warning about + collisions. A class defined later is warned about once agents look at types.""" + typeinfo.type_names() + + def audit(self): + """Log likely exposure mistakes once the app is set up.""" + if not self.exposes_anything(): + return + try: + catalog = self._catalog(self._graph()) + except Exception: + logger.exception("Invalid object exposure") + return + + permissions = self._mcp.permissions + for name, info in catalog.items(): + needs = {f"{name}.{m.attr}": m.requires for m in info.methods.values()} + for prop in info.properties.values(): + needs[f"{name}.{prop.name}"] = prop.requires + prop.set_requires + if info.constructor is not None: + needs[f"{name} constructor"] = info.constructor.requires + for what, requires in needs.items(): + unknown = [n for n in requires if n not in permissions] + if unknown: + logger.warning( + "%s requires undeclared permissions %s; it is never offered", + what, + unknown, + ) + for field_name, held in info.hidden_references.items(): + visible = sorted( + type_name(k) for t in held for k in subclasses(t) if is_visible(k) + ) + if visible: + logger.warning( + "%s.%s holds objects agents see (%s) but is not exposed, so " + "they are unreachable through it; expose it, or hide it " + "explicitly with mcp=False", + name, + field_name, + ", ".join(visible), + ) + + +def _objects_within(value, is_obj, limit=None): + """The objects in ``value``, at any depth of lists, tuples, sets and dict + values. ``limit`` keeps the first items of each list or tuple, as a + capped value shows them.""" + if is_obj(value): + yield value + elif isinstance(value, (list, tuple)): + for item in value[:limit]: + yield from _objects_within(item, is_obj, limit) + elif isinstance(value, (set, frozenset)): + for item in value: + yield from _objects_within(item, is_obj, limit) + elif isinstance(value, dict): + for item in value.values(): + yield from _objects_within(item, is_obj, limit) + + +def _holds_objects(value): + return next(_objects_within(value, is_object), None) is not None + + +# ----------------------------------------------------------------------------- +# object_types' index +# ----------------------------------------------------------------------------- + +# The end of a sentence: a period followed by a capital or the end (so not +# "e.g. a"), or by the paragraph's end +_SENTENCE_END = re.compile(r"[.!?](?=\s+[A-Z`(\"']|\s*$)") +MAX_ABOUT = 200 + + +def _first_sentence(text): + text = " ".join((text or "").split()) + match = _SENTENCE_END.search(text) + sentence = text[: match.end()] if match else text + if len(sentence) > MAX_ABOUT: + sentence = sentence[: MAX_ABOUT - 1] + "…" + return sentence + + +def _extending(entry, base_name, base): + """``entry`` as what it adds to, changes in or drops from ``base``. What it + lacks is False, including the constructor when it lost the base's.""" + delta = {"extends": base_name} + if "description" in entry: + delta["description"] = entry["description"] + for key in ("properties", "methods"): + delta[key] = {k: v for k, v in entry[key].items() if base[key].get(k) != v} + delta[key].update((k, False) for k in base[key] if k not in entry[key]) + if entry.get("constructor") != base.get("constructor"): + delta["constructor"] = entry.get("constructor", False) + return delta + + +def _index_entry(entry): + """A type as the index lists it: what it extends, the first sentence of its + description, its own property and method names, those of its base that it + lacks, and its constructor (False when it lost the base's).""" + index, lacks = {}, [] + about = _first_sentence(entry.get("description")) + if about: + index["about"] = about + for key in ("properties", "methods"): + names = [k for k, v in entry.get(key, {}).items() if v is not False] + lacks.extend(k for k, v in entry.get(key, {}).items() if v is False) + if names: + index[key] = names + if lacks: + index["lacks"] = lacks + if "constructor" in entry: + index["constructor"] = entry["constructor"] is not False + if "extends" in entry: + index["extends"] = entry["extends"] + return index + + +# ----------------------------------------------------------------------------- +# object_batch's results +# ----------------------------------------------------------------------------- + + +def _batch_failure(index, step_type, step, error, results): + if not isinstance(error, (ValueError, LookupError, PermissionError)): + logger.error("object_batch step %d raised", index, exc_info=error) + reason = error_text(error) + if step_type == "new": + what = f"new {step['new']}" + elif step_type == "call": + what = f"call {step['method']} on {step['call']}" + else: + what = f"update {step['update']}" + message = f"steps[{index}] ({what}) failed: {reason}" + if not results: # the first step: its own error says what it did + return message + if not message.endswith("."): + message += "." + applied = json.dumps(results, default=repr) + if index == 1: + return f"{message} Step 0 stays applied (no rollback); its result: {applied}" + return ( + f"{message} Steps 0-{index - 1} stay applied (no rollback); " + f"their results: {applied}" + ) + + +def _is_reference(value) -> bool: + return isinstance(value, dict) and "id" in value and value.keys() <= {"id", "type"} + + +def _holds_references(value) -> bool: + """Whether a rendered value is a reference or a list of them. Those are the + values an object keeps with ``results: "links"``.""" + if _is_reference(value): + return True + return isinstance(value, list) and bool(value) and all(map(_is_reference, value)) diff --git a/src/trame_server/mcp/permissions.py b/src/trame_server/mcp/permissions.py new file mode 100644 index 0000000..18d3e4e --- /dev/null +++ b/src/trame_server/mcp/permissions.py @@ -0,0 +1,275 @@ +"""Permissions: capabilities that the people running an application can +withhold from agents (``mcp.permissions``, ``--mcp-deny``). + +The application declares each capability it controls, and names the +permissions that a tool, a method, a property or a constructor needs with +``requires=``: + +.. code-block:: python + + server.mcp.permissions.declare("files", "Read files and folders on this computer") + server.mcp.permissions.declare( + "scripts", "Run Python code the agent writes", requires="files" + ) + + + @mcp.method(requires="files") + def list_files(self, path: str) -> dict: ... + +While a permission is denied, what needs it is left out of what agents are +offered (``tools/list``, ``object_types``), and calls to it are refused +with :class:`PermissionDeniedError`. A client that listed the tools +earlier still has them, and nothing tells it the list changed. +``require()`` checks inside a call, for a need that depends on the +arguments. + +``requires`` on a permission means what it means on a tool: the +permission is allowed only while the ones it requires are, typically +because it grants them anyway (a script can read files). Allowing it while +one of them stays denied raises, and denying one denies the ones requiring +it, so no setting holds one without the other. + +``--mcp-deny NAME[,NAME]`` (repeatable) and ``TRAME_MCP_DENY`` deny +permissions for the whole run. Whoever launched the application decided, +and neither its code nor its UI can allow them again (``locked`` tells a +UI which switches to disable). A name the application does not declare +refuses every agent call once the server runs: the user meant to withhold +something, and a typo must not leave it allowed. +""" + +from __future__ import annotations + +import logging +from collections.abc import Mapping +from dataclasses import dataclass + +from trame_common.obj.mcp import permission_names + +logger = logging.getLogger("trame_server.mcp") + +__all__ = ["DENY_ENV", "PermissionDeniedError", "Permissions"] + +DENY_ENV = "TRAME_MCP_DENY" + + +class PermissionDeniedError(PermissionError): + """A call needs a permission that is denied (``mcp.permissions``).""" + + +def launch_names(values) -> tuple[str, ...]: + """Names given at launch, each value comma-separated.""" + return tuple(n.strip() for v in values if v for n in v.split(",") if n.strip()) + + +@dataclass +class _Permission: + description: str + requires: tuple[str, ...] + allowed: bool + + +class Permissions(Mapping): + """The declared permissions (``mcp.permissions``), as a mapping from each + name to whether agents are allowed it. Assign ``permissions[name] = + bool``, or call :meth:`update` to set several at once, at any time.""" + + def __init__(self): + self._declared: dict[str, _Permission] = {} + self._launch_denied = frozenset() + self._started = False + + def __getitem__(self, name) -> bool: + return self._get(name).allowed + + def __iter__(self): + return iter(self._declared) + + def __len__(self): + return len(self._declared) + + def __setitem__(self, name, allowed): + self.update({name: allowed}) + + def __repr__(self): + return f"Permissions({dict(self)!r})" + + def declare(self, name, description, *, requires=None, allowed=True): + """Declare a permission, allowed unless ``allowed=False``. + + :param description: what it lets an agent do, in a few words. An agent + reads it when a call is refused. + :param requires: the permissions, declared earlier, that must be allowed + for this one to be, such as those it grants anyway. It starts denied + when one of them is denied. + """ + if not isinstance(name, str) or not name: + msg = f"A permission name is a non-empty string, got {name!r}" + raise TypeError(msg) + if name in self._declared: + msg = f"The permission {name!r} is already declared" + raise ValueError(msg) + requires = permission_names(requires) + unknown = [n for n in requires if n not in self._declared] + if unknown: + msg = f"{name!r} requires {unknown}, which are not declared (declare them first)" + raise ValueError(msg) + allowed = ( + bool(allowed) + and name not in self._launch_denied + and all(self._declared[n].allowed for n in requires) + ) + self._declared[name] = _Permission(str(description), requires, allowed) + + def update(self, values=(), /, **kwargs): + """Allow (True) or deny (False) several permissions at once, all or + none. + + A permission is allowed only with those it requires: allowing one + while one of them stays denied raises ``ValueError``, so allow both + in the same call. Denying one also denies those requiring it. A + permission denied at launch (``locked``) stays denied. + """ + changes = {**dict(values), **kwargs} + for name in changes: + self._get(name) + deny = self._closure({n for n, v in changes.items() if not v}, self._requiring) + # Locked ones are left out (what requires a locked one is locked too) + allowing = {n for n, v in changes.items() if v} + held = allowing & set(self.locked) + allowed = allowing - held + + def ends_allowed(name): + return name in allowed or ( + self._declared[name].allowed and name not in deny + ) + + missing = {} + for name in self._ordered(allowed): + needed = self._closure(self._declared[name].requires, self._required) + needed = [r for r in self._ordered(needed) if not ends_allowed(r)] + if needed: + missing[name] = needed + if missing: + needs = "; ".join( + f"{name!r} requires {', '.join(map(repr, needed))}" + for name, needed in missing.items() + ) + msg = ( + f"Cannot allow what requires a denied permission ({needs}); allow " + f"them in the same update" + ) + raise ValueError(msg) + if held: + logger.info( + "Not allowing %s: denied at launch for this run", + ", ".join(sorted(held)), + ) + for name in allowed: + self._declared[name].allowed = True + for name in deny: + self._declared[name].allowed = False + + @property + def locked(self) -> tuple[str, ...]: + """The permissions denied at launch, and those requiring them: + denied for the whole run, whatever is assigned.""" + held = self._closure( + self._launch_denied & self._declared.keys(), self._requiring + ) + return self._ordered(held) + + @property + def denied(self) -> tuple[str, ...]: + """The names of the denied permissions, in declaration order.""" + return tuple(name for name, p in self._declared.items() if not p.allowed) + + def description(self, name) -> str: + return self._get(name).description + + def allows(self, names) -> bool: + """Whether every permission in ``names`` is declared and allowed.""" + return all( + name in self._declared and self._declared[name].allowed for name in names + ) + + def require(self, *names, what="This call"): + """Raise :class:`PermissionDeniedError` unless every permission in ``names`` is + allowed. ``what`` names the refused action.""" + for name in names: + permission = self._declared.get(name) + if permission is None: + msg = f'{what}: permission "{name}" was never declared by the app' + + raise PermissionDeniedError(msg) + if not permission.allowed: + msg = ( + f'{what}: permission "{name}" ({permission.description}) is off; ' + f"only the people running the app can turn it on." + ) + raise PermissionDeniedError(msg) + + # ------------------------------------------------------------------------- + # Launch options and lifecycle (driven by the MCP) + # ------------------------------------------------------------------------- + + def _deny_at_launch(self, names): + self._launch_denied = frozenset(names) + + def _start(self): + """Called when the server starts: report the launch names that match no + permission. From then on, they refuse every agent call.""" + self._started = True + unknown = self._unknown_at_launch() + if unknown: + logger.error( + "--mcp-deny names %s, which this application does not declare " + "(it declares: %s); every agent call is refused", + ", ".join(unknown), + ", ".join(self._declared) or "none", + ) + + def _unknown_at_launch(self) -> list[str]: + return sorted(self._launch_denied - self._declared.keys()) + + def _check_launch(self): + """Refuse agent calls while a launch name matches no permission.""" + unknown = self._started and self._unknown_at_launch() + if unknown: + msg = ( + f"Agent access is refused: the application was launched denying " + f"{', '.join(unknown)}, which it does not declare" + ) + raise PermissionDeniedError(msg) + + # ------------------------------------------------------------------------- + # Helpers + # ------------------------------------------------------------------------- + + def _get(self, name) -> _Permission: + try: + return self._declared[name] + except (KeyError, TypeError): + msg = ( + f"No permission {name!r} is declared; declared: {list(self._declared)}" + ) + raise KeyError(msg) from None + + def _ordered(self, names) -> tuple[str, ...]: + """``names`` in declaration order.""" + return tuple(name for name in self._declared if name in names) + + def _required(self, name): + return self._declared[name].requires + + def _requiring(self, name): + return [n for n, p in self._declared.items() if name in p.requires] + + @staticmethod + def _closure(names, step) -> set[str]: + found, stack = set(), list(names) + while stack: + name = stack.pop() + if name not in found: + found.add(name) + stack.extend(step(name)) + return found diff --git a/src/trame_server/mcp/progress.py b/src/trame_server/mcp/progress.py new file mode 100644 index 0000000..e9e7439 --- /dev/null +++ b/src/trame_server/mcp/progress.py @@ -0,0 +1,164 @@ +"""Progress notifications: what a tool or method reports through its +``Progress`` parameter, sent to the agent that asked for it. + +A request carrying a ``progressToken`` in its ``_meta`` gets a +:class:`Reporter` for its duration, when its transport can send +notifications (``MCPProtocol.handle``'s ``notify``). The request's calls find +it in a context variable, which follows them into the tasks they create. +:func:`inject` passes it to a function that declares a ``Progress`` +parameter, or passes a reporter that does nothing when the agent asked for no +progress. + +MCP requires the value sent to grow, and nothing to be sent after the result. +So a report whose value doesn't grow is dropped, and its message goes with the +next one. Reports are also combined into at most one notification per +``MIN_INTERVAL``, the latest winning. The last one still waiting when the +call returns goes out just before the result. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import contextvars +import logging + +from trame_common.obj.mcp import Progress + +from trame_server.mcp.schema import progress_parameter + +logger = logging.getLogger("trame_server.mcp") + +__all__ = ["Reporter", "inject", "reporting"] + +# Seconds between two notifications of one request, at least +MIN_INTERVAL = 0.1 + +# The decimal places by which a batch step's report without a total moves the +# value: about half a float's precision, leaving the other half for the steps +STEP_DIGITS = 8 + +_NO_PROGRESS = Progress() +_current = contextvars.ContextVar("trame_mcp_progress", default=None) + + +class Reporter(Progress): + """Sends a request's progress with ``send(fields)``, an async callable + taking the notification's ``progress``, ``total`` and ``message``.""" + + def __init__(self, send, interval=MIN_INTERVAL): + self._send = send + self._interval = interval + self._sent = None # the last value sent + self._pending = None # (value, total) not sent yet + self._message = None # goes with the next notification + self._next = 0.0 # loop time from which the next one may go + self._flush = None # the task sending what is pending, when due + self._sending = asyncio.Lock() + self._closed = False + + async def __call__(self, value, total=None, message=None): + if self._closed: + return + if message is not None: + self._message = message + latest = self._pending[0] if self._pending else self._sent + if latest is not None and not value > latest: + return + self._pending = (value, total) + if self._flush is not None: + return # going out when due, as the latest + delay = self._next - asyncio.get_running_loop().time() + if delay <= 0: + await self._send_pending() + else: + self._flush = asyncio.ensure_future(self._send_later(delay)) + + async def _send_later(self, delay): + await asyncio.sleep(delay) + self._flush = None + # Shielded: closing waits for a notification being written + await asyncio.shield(self._send_pending()) + + async def _send_pending(self): + async with self._sending: + if self._closed or self._pending is None: + return + value, total = self._pending + fields = {"progress": value} + if total is not None: + fields["total"] = total + if self._message is not None: + fields["message"] = self._message + self._pending, self._message, self._sent = None, None, value + self._next = asyncio.get_running_loop().time() + self._interval + try: + await self._send(fields) + except Exception: # the client went away: the call goes on + logger.debug("Sending progress failed", exc_info=True) + self._closed = True + + async def close(self): + """Called when the result is ready: send what is pending now, before the + result, and nothing after it.""" + if self._flush is not None: + self._flush.cancel() + self._flush = None + await self._send_pending() + async with self._sending: + self._closed = True + + +class _BatchStep(Progress): + """The reporter a method gets during one step of an ``object_batch``. Its + reports become progress within that step of the batch (step k of n). + With a total, the value moves by the method's fraction of the step; + without one, each report moves it by ``10**-STEP_DIGITS``, so that it + still grows and goes out with its message.""" + + def __init__(self, batch, index, count): + self._batch = batch + self._index = index + self._count = count + self._done = index # the furthest within this step + + async def __call__(self, value, total=None, message=None): + if total: + done = self._index + min(max(value / total, 0.0), 1.0) + else: + tick = 10**-STEP_DIGITS + done = min(round(self._done + tick, STEP_DIGITS), self._index + 1 - tick) + self._done = max(self._done, done) + await self._batch(done, total=self._count, message=message) + + +@contextlib.contextmanager +def reporting(reporter): + """Make ``reporter`` the reporter of the calls made in the block (None when + the agent asked for no progress).""" + token = _current.set(reporter) + try: + yield + finally: + _current.reset(token) + + +def batch_step(index, count): + """The reporter that the method of step ``index`` of ``count`` gets in a batch, + or None when the agent asked for no progress.""" + batch = _current.get() + return None if batch is None else _BatchStep(batch, index, count) + + +def current(): + """The reporter of the request being handled, or None.""" + return _current.get() + + +def inject(fn, arguments): + """``arguments``, with the reporter added when ``fn`` declares a + ``Progress`` parameter: the request's, or one that does nothing.""" + name = progress_parameter(fn) + if name is None: + return arguments + return {**arguments, name: _current.get() or _NO_PROGRESS} diff --git a/src/trame_server/mcp/protocol.py b/src/trame_server/mcp/protocol.py new file mode 100644 index 0000000..4b722ba --- /dev/null +++ b/src/trame_server/mcp/protocol.py @@ -0,0 +1,560 @@ +"""MCP JSON-RPC method dispatch (transport-independent). + +Clients of both eras are served, as MCP 2026-07-28 lets a server do. A +request that names its protocol version in ``_meta`` comes from a client of +2026-07-28 or later: it carries everything needed to serve it, and its +result says so (``resultType``, the server's identity, caching hints). A +client that opens with ``initialize`` speaks an earlier revision and is +served as that revision says. +""" + +from __future__ import annotations + +import asyncio +import base64 +import contextlib +import inspect +import json +import logging +import re + +from trame_server.mcp import call_log, progress +from trame_server.mcp.content import to_error_result, to_tool_result +from trame_server.mcp.schema import ToolError, ValidationError, bind_arguments + +logger = logging.getLogger("trame_server.mcp") + +__all__ = ["MCPProtocol"] + +# 2026-07-28 and later: every request names its version and the client's +# capabilities in _meta, and there is no handshake +MODERN_VERSIONS = ("2026-07-28",) +# Earlier revisions open with an initialize handshake, which negotiates one +# of these (the first is preferred); still served for clients that have not +# moved on +LEGACY_VERSIONS = ("2025-11-25", "2025-06-18", "2025-03-26", "2024-11-05") +PROTOCOL_VERSIONS = MODERN_VERSIONS + LEGACY_VERSIONS + +# Reserved _meta keys +VERSION_KEY = "io.modelcontextprotocol/protocolVersion" +CAPABILITIES_KEY = "io.modelcontextprotocol/clientCapabilities" +CLIENT_INFO_KEY = "io.modelcontextprotocol/clientInfo" +SERVER_INFO_KEY = "io.modelcontextprotocol/serverInfo" + +# The results that carry caching hints (2026-07-28) +CACHEABLE = frozenset( + ( + "server/discover", + "tools/list", + "prompts/list", + "resources/list", + "resources/templates/list", + "resources/read", + ) +) + +# JSON-RPC error codes +PARSE_ERROR = -32700 +INVALID_REQUEST = -32600 +METHOD_NOT_FOUND = -32601 +INVALID_PARAMS = -32602 +INTERNAL_ERROR = -32603 +# MCP's own (2026-07-28) +HEADER_MISMATCH = -32020 +UNSUPPORTED_VERSION = -32022 +# The reply to a request the client cancelled (LSP's code; MCP defines none). +# The client ignores it, but an HTTP POST has to end with something. +REQUEST_CANCELLED = -32800 + + +def _error(rpc_id, code, message, data=None): + error = {"code": code, "message": message} + if data is not None: + error["data"] = data + return {"jsonrpc": "2.0", "id": rpc_id, "error": error} + + +def _result(rpc_id, payload): + return {"jsonrpc": "2.0", "id": rpc_id, "result": payload} + + +async def _call(fn, arguments): + value = fn(**arguments) + if inspect.isawaitable(value): + value = await value + return value + + +def _meta(params): + """A request's ``_meta``, or an empty dict.""" + meta = params.get("_meta") if isinstance(params, dict) else None + return meta if isinstance(meta, dict) else {} + + +def request_version(message): + """The protocol version a request names in its ``_meta``, as clients of + 2026-07-28 and later do; None for a client of an earlier revision.""" + return _meta(message.get("params")).get(VERSION_KEY) + + +def _caller(message, session): + """Who is asking, for the call log: the start of an HTTP session id, + else the name the client gives, kept to a few safe characters.""" + if isinstance(session, str): + return session[:6] + info = _meta(message.get("params")).get(CLIENT_INFO_KEY) + name = info.get("name") if isinstance(info, dict) else None + if not isinstance(name, str): + return None + return re.sub(r"[^A-Za-z0-9_.-]", "_", name)[:24] or None + + +def _reporter(params, notify): + """The progress reporter of a request asking for progress, when its + transport can notify; None otherwise.""" + token = _meta(params).get("progressToken") + if notify is None or isinstance(token, bool) or not isinstance(token, (str, int)): + return None + + async def send(fields): + await notify( + { + "jsonrpc": "2.0", + "method": "notifications/progress", + "params": {"progressToken": token, **fields}, + } + ) + + return progress.Reporter(send) + + +def check_request(message): + """The error that a request naming its version gets before it is handled (an + unsupported version, or missing client capabilities), or None. The HTTP + endpoint answers these errors with 400.""" + rpc_id = message.get("id") + version = request_version(message) + if version not in MODERN_VERSIONS: + hint = " (it opens with initialize)" if version in LEGACY_VERSIONS else "" + return _error( + rpc_id, + UNSUPPORTED_VERSION, + f"Unsupported protocol version: {version!r}{hint}", + {"supported": list(PROTOCOL_VERSIONS), "requested": version}, + ) + if not isinstance(message["params"]["_meta"].get(CAPABILITIES_KEY), dict): + return _error(rpc_id, INVALID_PARAMS, f"_meta must carry {CAPABILITIES_KEY}") + return None + + +def _describe_failure(name, error): + """Log a tool's exception, and return the text the agent reads. Expected + failures are logged at DEBUG, the others as errors with their traceback.""" + if isinstance(error, (ValueError, LookupError, PermissionError)): + # Expected, actionable failures (bad input, unknown target, ...) + logger.debug("Tool %s failed: %s", name, error) + if isinstance(error, (ToolError, PermissionError)): + return str(error) + else: + logger.error("Tool %s raised", name, exc_info=error) + return f"{type(error).__name__}: {error}" + + +class MCPProtocol: + """Dispatches MCP JSON-RPC messages against an ``MCP`` instance's registries.""" + + def __init__(self, mcp): + self._mcp = mcp + self._call_lock = asyncio.Lock() + # (session, request id) -> task of a request still running, and the + # tasks cancelled by the client -> its reason + self._in_flight = {} + self._cancelled = {} + served = { + "tools/list": self._on_tools_list, + "tools/call": self._on_tools_call, + "resources/list": self._on_resources_list, + "resources/templates/list": self._on_resource_templates_list, + "resources/read": self._on_resources_read, + "prompts/list": self._on_prompts_list, + "prompts/get": self._on_prompts_get, + } + self._modern_handlers = {"server/discover": self._on_discover, **served} + self._legacy_handlers = { + "initialize": self._on_initialize, + "ping": self._on_ping, + **served, + } + + async def handle(self, message: dict, *, session=None, notify=None) -> dict | None: + """Handle one JSON-RPC message. Returns a response dict, or None for + notifications and responses, which expect no reply. + + ``session`` scopes request ids, so that ``notifications/cancelled`` only + reaches the requests of the client that sent it. The HTTP endpoint passes + the ``Mcp-Session-Id`` header of clients older than 2026-07-28, and an + in-process client passes itself. A request the client cancels is answered + with a ``REQUEST_CANCELLED`` error, which the client ignores. + + ``notify``, an async callable, sends the client a notification about this + request before the response; the transport passes one when it can send + notifications. With it, a request asking for progress (a + ``progressToken`` in its ``_meta``) gets the progress its tool or method + reports. + """ + if not isinstance(message, dict) or message.get("jsonrpc") != "2.0": + return _error(None, INVALID_REQUEST, "Invalid JSON-RPC 2.0 message") + + method = message.get("method") + rpc_id = message.get("id") + + if method is None: + # A response from the client to a server request needs nothing more. + return None + if not isinstance(method, str): + return _error(rpc_id, INVALID_REQUEST, "Invalid method") + + # The calls made while handling it know who asked (the call log) + with call_log.request(_caller(message, session)): + return await self._dispatch(message, method, rpc_id, session, notify) + + async def _dispatch(self, message, method, rpc_id, session, notify): + if method.startswith("notifications/"): + if method == "notifications/cancelled": + params = message.get("params") + if isinstance(params, dict): + self.cancel( + params.get("requestId"), + session=session, + reason=params.get("reason"), + ) + return None + + modern = request_version(message) is not None + if modern: + rejection = check_request(message) + if rejection is not None: + return rejection + handler = self._handler(method, modern) + if handler is None: + return _error(rpc_id, METHOD_NOT_FOUND, f"Method not found: {method}") + + params = message.get("params") or {} + if not isinstance(params, dict): + return _error(rpc_id, INVALID_PARAMS, "params must be an object") + + reporter = _reporter(params, notify) + try: + with progress.reporting(reporter): + payload = handler(params) + if inspect.isawaitable(payload): + payload = await self._run_cancellable(payload, (session, rpc_id)) + except _RequestCancelled as e: + return _error(rpc_id, REQUEST_CANCELLED, str(e)) + except _InvalidParams as e: + return _error(rpc_id, INVALID_PARAMS, str(e)) + except PermissionError as e: + return _error(rpc_id, INVALID_REQUEST, str(e)) + except Exception as e: + logger.exception("Internal error handling %s", method) + return _error(rpc_id, INTERNAL_ERROR, f"{type(e).__name__}: {e}") + finally: + if reporter is not None: + await reporter.close() # nothing is sent after the response + + if modern: + payload = self._complete(method, payload) + return _result(rpc_id, payload) + + def _handler(self, method, modern): + handlers = self._modern_handlers if modern else self._legacy_handlers + return handlers.get(method) + + def handles(self, message) -> bool: + """Whether a request's method is served, in its client's revision.""" + return ( + self._handler(message.get("method"), request_version(message) is not None) + is not None + ) + + def _complete(self, method, payload): + """A result as 2026-07-28 shapes it.""" + result = {"resultType": "complete", **payload} + result["_meta"] = { + **(payload.get("_meta") or {}), + SERVER_INFO_KEY: self._info(), + } + if method in CACHEABLE: + # What it lists changes as the app runs, with no notification + # sent, and describes one app's live state + result["ttlMs"] = 0 + result["cacheScope"] = "private" + return result + + def _info(self): + from trame_server import __version__ # noqa: PLC0415 + + return {"name": self._mcp.name, "version": __version__} + + # ------------------------------------------------------------------------- + # Cancellation + # ------------------------------------------------------------------------- + + def cancel(self, request_id, *, session=None, reason=None) -> bool: + """Cancel the running request ``request_id`` of ``session``, as + ``notifications/cancelled`` does. Returns whether one was running. + + The request's task is cancelled: its tool stops at its next + ``await`` (a synchronous tool has finished before any cancel can + arrive), releasing the call lock and flushing the state it changed. + """ + try: + task = self._in_flight.get((session, request_id)) + except TypeError: # an unhashable id: no request has it + return False + if task is None or task.done(): + return False + logger.info( + "Request %r cancelled by the client: %s", request_id, reason or "no reason" + ) + self._cancelled[task] = reason + task.cancel() + return True + + async def _run_cancellable(self, awaitable, key): + """Await ``awaitable`` as a task that ``cancel(key)`` can stop.""" + task = asyncio.ensure_future(awaitable) + try: + registered = self._in_flight.setdefault(key, task) is task + except TypeError: + registered = False + try: + return await task + except asyncio.CancelledError: + current = asyncio.current_task() + # Task.cancelling() is Python 3.11+; without it, assume not + outer_cancelled = ( + current is not None and getattr(current, "cancelling", int)() + ) + if task not in self._cancelled or outer_cancelled: + raise # cancelled from the outside, e.g. the client disconnected + reason = self._cancelled[task] + msg = f"Request cancelled: {reason}" if reason else "Request cancelled" + raise _RequestCancelled(msg) from None + finally: + self._cancelled.pop(task, None) + if registered and self._in_flight.get(key) is task: + del self._in_flight[key] + + # ------------------------------------------------------------------------- + # Lifecycle + # ------------------------------------------------------------------------- + + def _on_discover(self, _params): + return { + "supportedVersions": list(PROTOCOL_VERSIONS), + "capabilities": {"tools": {}, "resources": {}, "prompts": {}}, + "instructions": self._mcp._agent_instructions(), + } + + def _on_initialize(self, params): + """The handshake of clients before 2026-07-28.""" + requested = params.get("protocolVersion") + version = requested if requested in LEGACY_VERSIONS else LEGACY_VERSIONS[0] + payload = { + "protocolVersion": version, + "capabilities": { + "tools": {"listChanged": False}, + "resources": {"listChanged": False}, + "prompts": {"listChanged": False}, + }, + "serverInfo": self._info(), + } + payload["instructions"] = self._mcp._agent_instructions() + return payload + + def _on_ping(self, _params): + return {} + + # ------------------------------------------------------------------------- + # Tools + # ------------------------------------------------------------------------- + + def _on_tools_list(self, _params): + tools = [] + for entry in self._mcp.list_tools(): + descriptor = { + "name": entry.name, + "description": entry.description, + "inputSchema": self._mcp._input_schema(entry), + } + if entry.annotations: + descriptor["annotations"] = entry.annotations + tools.append(descriptor) + return {"tools": tools} + + async def _on_tools_call(self, params): + name = params.get("name") + entry = self._mcp.get_tool(name) if isinstance(name, str) else None + if entry is None: + msg = f"Unknown tool: {name}" + with call_log.logged("tool", str(name), params.get("arguments")): + raise _InvalidParams(msg) + + try: + value = await self._execute(entry, params.get("arguments")) + except Exception as e: + return to_error_result(_describe_failure(name, e)) + + return to_tool_result(value, structured=self._mcp.structured_content) + + async def _execute(self, entry, arguments): + """Run a tool entry. Access, permissions and arguments are checked first, then + the call runs, one at a time unless it may run concurrently. + + Returns the tool's raw value; raises on failure. + """ + with call_log.logged("tool", entry.name, arguments): + self._mcp._check_enabled() + self._mcp.permissions.require(*entry.requires, what=entry.name) + try: + arguments = bind_arguments( + entry.fn, arguments, self._mcp._input_schema(entry) + ) + except ValidationError as e: + msg = f"Invalid arguments for tool {entry.name}: {e}" + raise ValidationError(msg) from None + + async with self.locked(entry.concurrent): + # Ids given for object-typed parameters become objects here, + # under the lock, where the graph is as earlier calls left it. + arguments = self._mcp._objects.resolve_arguments( + entry.fn, + arguments, + read_only=bool((entry.annotations or {}).get("readOnlyHint")), + ) + arguments = progress.inject(entry.fn, arguments) + result = await self._call_flushing(entry.fn, arguments) + if entry.input_schema is None: # the app's own tool + result = self._mcp._objects.encode_tool_result(result) + return result + + @contextlib.asynccontextmanager + async def locked(self, concurrent=False): + """Hold the call lock for the block, so that agent calls run one at a time. + Nothing is held when the call is ``concurrent``, or when every call may be + (``concurrent_calls``). + + ``object_call`` runs unlocked and takes the lock here itself, so the + called method's own ``concurrent`` option decides. + """ + if not (concurrent or self._mcp.concurrent_calls): + async with self._call_lock: + yield + else: + yield + + async def _call_flushing(self, fn, arguments): + # trame pushes state changes and runs @state.change reactions on a + # flush, which an HTTP request never triggers by itself: without one, + # what a tool assigns would wait for some unrelated flush. + with self._mcp.server.state: + return await _call(fn, arguments) + + # ------------------------------------------------------------------------- + # Resources + # ------------------------------------------------------------------------- + + def _on_resources_list(self, _params): + return { + "resources": [ + { + "uri": entry.uri, + "name": entry.name, + **({"description": entry.description} if entry.description else {}), + **({"mimeType": entry.mime_type} if entry.mime_type else {}), + } + for entry in self._mcp.list_resources() + ] + } + + def _on_resource_templates_list(self, _params): + return {"resourceTemplates": []} + + async def _on_resources_read(self, params): + uri = params.get("uri") + entry = self._mcp.get_resource(uri) if isinstance(uri, str) else None + if entry is None: + msg = f"Unknown resource: {uri}" + with call_log.logged("resource", str(uri)): + raise _InvalidParams(msg) + + with call_log.logged("resource", entry.uri): + self._mcp._check_enabled() + value = await _call(entry.fn, {}) + if isinstance(value, bytes): + content = { + "uri": uri, + "mimeType": entry.mime_type or "application/octet-stream", + "blob": base64.b64encode(value).decode("ascii"), + } + else: + if not isinstance(value, str): + value = json.dumps(value) + content = { + "uri": uri, + "mimeType": entry.mime_type or "text/plain", + "text": value, + } + return {"contents": [content]} + + # ------------------------------------------------------------------------- + # Prompts + # ------------------------------------------------------------------------- + + def _on_prompts_list(self, _params): + return { + "prompts": [ + { + "name": entry.name, + **({"description": entry.description} if entry.description else {}), + "arguments": entry.arguments, + } + for entry in self._mcp.list_prompts() + ] + } + + async def _on_prompts_get(self, params): + name = params.get("name") + entry = self._mcp.get_prompt(name) if isinstance(name, str) else None + if entry is None: + msg = f"Unknown prompt: {name}" + with call_log.logged("prompt", str(name), params.get("arguments")): + raise _InvalidParams(msg) + + arguments = params.get("arguments") or {} + with call_log.logged("prompt", name, arguments): + self._mcp._check_enabled() + try: + inspect.signature(entry.fn).bind(**arguments) + except TypeError as e: + msg = f"Invalid arguments for prompt {name}: {e}" + raise _InvalidParams(msg) from None + + value = await _call(entry.fn, arguments) + if isinstance(value, str): + messages = [{"role": "user", "content": {"type": "text", "text": value}}] + else: + messages = value + return { + **({"description": entry.description} if entry.description else {}), + "messages": messages, + } + + +class _InvalidParams(ValueError): + pass + + +class _RequestCancelled(Exception): + pass diff --git a/src/trame_server/mcp/schema.py b/src/trame_server/mcp/schema.py new file mode 100644 index 0000000..6c2d113 --- /dev/null +++ b/src/trame_server/mcp/schema.py @@ -0,0 +1,644 @@ +"""JSON Schema generation from Python type annotations, and validation of +incoming MCP values against the schemas we advertise.""" + +from __future__ import annotations + +import inspect +import json +import re +import types +import typing + +from trame_common.obj.mcp import MCPExposable, Progress + +__all__ = [ + "ToolError", + "ValidationError", + "bind_arguments", + "docstring", + "error_text", + "from_json", + "function_input_schema", + "is_object", + "object_class", + "parameter_hints", + "resolve_object_arguments", + "return_hint", + "type_to_schema", + "unwrap", + "validate", +] + +_SCALARS = { + str: {"type": "string"}, + int: {"type": "integer"}, + float: {"type": "number"}, + bool: {"type": "boolean"}, + type(None): {"type": "null"}, +} + + +def type_to_schema(annotation, custom=None) -> dict: + """Best-effort conversion of a type annotation to a JSON Schema fragment. + + Unknown or missing annotations yield ``{}`` (any type). ``custom`` is an + optional ``annotation -> schema | None`` hook consulted first at every + level (used to describe references to components and models). + + ``Annotated`` metadata contributes a description (``str``) or extra + schema keywords (``dict``), e.g. + ``Annotated[int, "Number of facets", {"minimum": 3, "maximum": 60}]``. + """ + if annotation is None or annotation is inspect.Parameter.empty: + return {} + + if custom is not None: + schema = custom(annotation) + if schema is not None: + return schema + + # isinstance first: Annotated metadata may be unhashable (dict schemas) + if isinstance(annotation, type) and annotation in _SCALARS: + return dict(_SCALARS[annotation]) + + origin = typing.get_origin(annotation) + args = typing.get_args(annotation) + + def convert(a): + return type_to_schema(a, custom) + + if origin is None: + if annotation is dict: + return {"type": "object"} + if annotation in (list, set, frozenset, tuple): + return {"type": "array"} + return {} + + if origin is typing.Annotated: + schema = convert(args[0]) + for extra in args[1:]: + if isinstance(extra, str): + schema["description"] = extra + elif isinstance(extra, dict): + schema.update(extra) + return schema + + if origin is typing.Literal: + return {"enum": list(args)} + + if origin in (typing.Union, types.UnionType): + non_none = [a for a in args if a is not type(None)] + optional = len(non_none) != len(args) + if len(non_none) == 1: + schema = convert(non_none[0]) + return nullable(schema) if optional else schema + variants = [convert(a) for a in non_none] + if optional: + variants.append({"type": "null"}) + return {"anyOf": variants} + + if origin is tuple and not (len(args) == 2 and args[1] is Ellipsis): + # Fixed-length tuple, e.g. tuple[float, float, float] + items = [convert(a) for a in args] + schema = {"type": "array", "minItems": len(items), "maxItems": len(items)} + if items and all(item == items[0] for item in items): + schema["items"] = items[0] + elif items: + schema["prefixItems"] = items + return schema + + if origin in (list, set, frozenset, tuple): + return {"type": "array", "items": convert(args[0]) if args else {}} + + if origin is dict: + value_schema = convert(args[1]) if len(args) == 2 else {} + return {"type": "object", "additionalProperties": value_schema} + + return {} + + +_ARRAYS = (list, tuple, set, frozenset) + + +def from_json(value, annotation, resolve=None, path="value"): + """ + ``value``, validated JSON, in the Python shape ``annotation`` declares. + JSON has no tuples or sets, so arrays become tuples or sets where the + annotation says so (``tuple[float, float, float]``, ``set[str]``), nested + in lists, dicts and unions too. Anything else is returned as is. + + With ``resolve``, an id or name given where the annotation names a class + whose instances agents reach becomes the object: ``resolve(id, cls, + path)`` returns it or raises. Converting a value twice gives the same + value, so ids can be resolved after the shape is restored. + """ + if value is None or annotation is None or annotation is inspect.Parameter.empty: + return value + if resolve is not None: + cls = object_class(annotation) + if cls is not None: + return resolve(value, cls, path) if isinstance(value, str) else value + origin = typing.get_origin(annotation) + args = typing.get_args(annotation) + if origin is typing.Annotated: + return from_json(value, args[0], resolve, path) + if origin in (typing.Union, types.UnionType): + for arg in args: + base = ( + typing.get_args(arg)[0] + if typing.get_origin(arg) is typing.Annotated + else arg + ) + container = typing.get_origin(base) or base + if isinstance(value, str): + if resolve is not None and object_class(base) is not None: + return resolve(value, base, path) + elif isinstance(value, dict): + if container is dict: + return from_json(value, base, resolve, path) + elif isinstance(value, _ARRAYS) and container in _ARRAYS: + return from_json(value, base, resolve, path) + return value + if isinstance(value, _ARRAYS): + if annotation in (tuple, set, frozenset): + return annotation(value) + if origin is tuple: + if len(args) == 2 and args[1] is Ellipsis: + args = (args[0],) * len(value) + if len(args) != len(value): + return tuple(value) + return tuple( + from_json(v, a, resolve, f"{path}[{i}]") + for i, (v, a) in enumerate(zip(value, args, strict=True)) + ) + if origin in (list, set, frozenset) and args: + items = [ + from_json(v, args[0], resolve, f"{path}[{i}]") + for i, v in enumerate(value) + ] + return items if origin is list else origin(items) + if isinstance(value, dict) and origin is dict and len(args) == 2: + return { + k: from_json(v, args[1], resolve, f"{path}.{k}") for k, v in value.items() + } + return value + + +def nullable(schema: dict) -> dict: + """Return ``schema`` extended to also accept ``null``.""" + if not schema or _accepts_null(schema): + return schema + return {"anyOf": [schema, {"type": "null"}]} + + +def _accepts_null(schema): + json_type = schema.get("type") + if json_type == "null" or (isinstance(json_type, list) and "null" in json_type): + return True + return any(_accepts_null(s) for s in schema.get("anyOf", ())) + + +def parameter_hints(fn): + """The resolved annotations of ``fn``'s parameters. One annotation that + cannot be resolved (a forward reference to a later class, typically the + return type) does not lose the others.""" + try: + return typing.get_type_hints(fn, include_extras=True) + except Exception: + pass + func = inspect.unwrap(getattr(fn, "__func__", fn)) + globalns = getattr(func, "__globals__", {}) + hints = {} + for name, annotation in getattr(func, "__annotations__", {}).items(): + if name == "return": + continue + holder = types.SimpleNamespace(__annotations__={name: annotation}) + try: + hints[name] = typing.get_type_hints( + holder, globalns=globalns, include_extras=True + )[name] + except Exception: + continue + return hints + + +def is_progress(annotation) -> bool: + """True for ``Progress``, ``Progress | None`` and ``Optional[Progress]`` + (Annotated or not): the parameter trame fills in with a reporter.""" + if typing.get_origin(annotation) is typing.Annotated: + annotation = typing.get_args(annotation)[0] + if typing.get_origin(annotation) in (typing.Union, types.UnionType): + args = [a for a in typing.get_args(annotation) if a is not type(None)] + annotation = args[0] if len(args) == 1 else None + return isinstance(annotation, type) and issubclass(annotation, Progress) + + +def progress_parameter(fn): + """The name of ``fn``'s ``Progress`` parameter, or None.""" + for name, annotation in parameter_hints(fn).items(): + if name != "return" and is_progress(annotation): + return name + return None + + +def check_progress(fn, label): + """Raise when ``fn`` declares a ``Progress`` parameter but isn't async, since + reports are awaited.""" + if progress_parameter(fn) is None or inspect.iscoroutinefunction(fn): + return + msg = f"{label}: a Progress parameter needs an async function (async def)" + raise TypeError(msg) + + +def _json_safe(value): + try: + json.dumps(value) + except (TypeError, ValueError): + return False + return True + + +def return_hint(fn): + """The resolved return annotation of ``fn``, or None.""" + try: + return typing.get_type_hints(fn, include_extras=True).get("return") + except Exception: + return None + + +# A line starting a list item or a field (":param x:"), which keeps its line +# break; the lines indented under it continue it +_ITEM = re.compile(r"\s*(?:[-*+•]|\d+[.)]|:[^:\s][^:]*:)\s") + + +def unwrap(text): + """``text`` with the wrapped lines of each paragraph joined, as agents read + it best. Blank lines, list items, fields and indented blocks (code, + ``Args:`` sections) keep their line breaks, and so does what is between + ``` fences.""" + if not text: + return text + lines = [] + fenced = False + # What the next line may continue: top-level prose, or the item at + # this indent (by lines indented under it) + prose, item = False, None + for line in text.splitlines(): + stripped = line.strip() + indent = len(line) - len(line.lstrip()) + fence = stripped.startswith("```") + starts_item = _ITEM.match(line) is not None + if stripped and not (fenced or fence or starts_item): + if (item is not None and indent > item) or (prose and indent == 0): + lines[-1] = f"{lines[-1].rstrip()} {stripped}" + continue + lines.append(line) + fenced ^= fence + if fenced or fence or not stripped: + prose, item = False, None + elif starts_item: + prose, item = False, indent + else: + prose, item = indent == 0, None + return "\n".join(lines) + + +def docstring(obj): + """``obj``'s docstring as agents read it (see :func:`unwrap`), or None.""" + return unwrap(inspect.getdoc(obj)) + + +def function_input_schema(fn, skip=("self",), ref=None) -> dict: + """Build the MCP ``inputSchema`` for a callable from its signature. + ``ref`` describes parameters annotated with a component or model class, + which take the object's id or name (see ``type_to_schema``'s ``custom``). A + ``Progress`` parameter is trame's to fill in, so agents don't see it.""" + sig = inspect.signature(fn) + hints = parameter_hints(fn) + + properties = {} + required = [] + var_keyword = False + for pname, param in sig.parameters.items(): + if param.kind is inspect.Parameter.VAR_KEYWORD: + var_keyword = True + continue + if pname in skip or param.kind is inspect.Parameter.VAR_POSITIONAL: + continue + if is_progress(hints.get(pname)): + continue + prop = type_to_schema(hints.get(pname), ref) + if param.default is inspect.Parameter.empty: + required.append(pname) + else: + if param.default is None: + prop = nullable(prop) + if _json_safe(param.default): + prop["default"] = param.default + properties[pname] = prop + + schema = {"type": "object", "properties": properties} + if required: + schema["required"] = required + if not var_keyword: + schema["additionalProperties"] = False + return schema + + +# ----------------------------------------------------------------------------- +# Validation +# ----------------------------------------------------------------------------- + + +class ToolError(ValueError): + """An expected, actionable failure whose message is complete as is: the + agent reads it without the exception's type name.""" + + +class ValidationError(ToolError): + """An MCP-supplied value does not match the advertised schema.""" + + +def error_text(error): + """The text agents read for a failure: the message of a ``ToolError`` or + ``PermissionError`` as is, else the exception's type and message.""" + if isinstance(error, (ToolError, PermissionError)): + return str(error) + return f"{type(error).__name__}: {error}" + + +def _type_name(value): + if value is None: + return "null" + if isinstance(value, bool): + return "boolean" + if isinstance(value, int): + return "integer" + if isinstance(value, float): + return "number" + if isinstance(value, str): + return "string" + if isinstance(value, (list, tuple)): + return "array" + if isinstance(value, dict): + return "object" + return type(value).__name__ + + +def _show(value): + """``value`` as JSON for a message, the way the agent wrote it.""" + try: + return json.dumps(value) + except (TypeError, ValueError): + return repr(value) + + +def _json_equal(a, b): + return a == b and isinstance(a, bool) == isinstance(b, bool) + + +def _check_type(value, json_type, path): + """Return ``value`` coerced to ``json_type`` (lossless numeric coercion only).""" + is_number = isinstance(value, (int, float)) and not isinstance(value, bool) + if json_type == "integer": + if is_number and isinstance(value, int): + return value + if is_number and isinstance(value, float) and value.is_integer(): + return int(value) + elif json_type == "number": + if is_number: + return float(value) if isinstance(value, int) else value + elif json_type == "string": + if isinstance(value, str): + return value + elif json_type == "boolean": + if isinstance(value, bool): + return value + elif json_type == "null": + if value is None: + return value + elif json_type == "array": + if isinstance(value, (list, tuple)): + return list(value) + elif json_type == "object": + if isinstance(value, dict): + return value + else: + return value # unknown type keyword: do not reject + + msg = f"{path}: expected {json_type}, got {_type_name(value)}" + raise ValidationError(msg) + + +def validate(value, schema, path="value"): + """Validate ``value`` against the JSON Schema subset trame generates. + + Returns the value with lossless coercions applied (``3.0`` for an + integer becomes ``3``, ``3`` for a number becomes ``3.0``); raises + :class:`ValidationError` naming the offending path otherwise. Keywords + outside the supported subset are ignored. + """ + if not schema: + return value + + branches = schema.get("anyOf") or schema.get("oneOf") + if branches: + errors = [] + # A branch of the value's own type first: 3 for float | int stays 3 + exact = _type_name(value) + branches = sorted(branches, key=lambda b: b.get("type") != exact) + for branch in branches: + try: + value = validate(value, branch, path) + break + except ValidationError as e: + errors.append(str(e).removeprefix(f"{path}: ")) + else: + msg = f"{path}: does not match any allowed form ({'; '.join(errors)})" + raise ValidationError(msg) + + if "object" in schema and not isinstance(value, str): + msg = ( + f"{path}: expected the id of an object of type {schema['object']}, " + f"got {_type_name(value)}" + ) + raise ValidationError(msg) + + json_type = schema.get("type") + if isinstance(json_type, list): + for candidate in sorted(json_type, key=lambda k: k != _type_name(value)): + try: + value = _check_type(value, candidate, path) + break + except ValidationError: + continue + else: + msg = f"{path}: expected {' or '.join(json_type)}, got {_type_name(value)}" + raise ValidationError(msg) + elif json_type is not None: + value = _check_type(value, json_type, path) + + if "const" in schema and not _json_equal(value, schema["const"]): + msg = f"{path}: must be {_show(schema['const'])}" + raise ValidationError(msg) + if "enum" in schema and not any(_json_equal(value, e) for e in schema["enum"]): + choices = ", ".join(_show(e) for e in schema["enum"]) + msg = f"{path}: must be one of {choices}, got {_show(value)}" + raise ValidationError(msg) + + if isinstance(value, (int, float)) and not isinstance(value, bool): + _check_bounds(value, schema, path) + elif isinstance(value, str): + _check_string(value, schema, path) + elif isinstance(value, list): + value = _check_array(value, schema, path) + elif isinstance(value, dict): + value = _check_object(value, schema, path) + + return value + + +def _check_bounds(value, schema, path): + checks = ( + ("minimum", lambda v, b: v >= b, "at least"), + ("maximum", lambda v, b: v <= b, "at most"), + ("exclusiveMinimum", lambda v, b: v > b, "greater than"), + ("exclusiveMaximum", lambda v, b: v < b, "less than"), + ) + for keyword, ok, label in checks: + if keyword in schema and not ok(value, schema[keyword]): + msg = f"{path}: must be {label} {schema[keyword]}, got {value}" + raise ValidationError(msg) + + +def _check_string(value, schema, path): + if "minLength" in schema and len(value) < schema["minLength"]: + msg = f"{path}: must be at least {schema['minLength']} characters" + raise ValidationError(msg) + if "maxLength" in schema and len(value) > schema["maxLength"]: + msg = f"{path}: must be at most {schema['maxLength']} characters" + raise ValidationError(msg) + if "pattern" in schema and not re.search(schema["pattern"], value): + msg = f"{path}: must match pattern {schema['pattern']!r}" + raise ValidationError(msg) + + +def _check_array(value, schema, path): + if "minItems" in schema and len(value) < schema["minItems"]: + msg = f"{path}: expected at least {schema['minItems']} items, got {len(value)}" + raise ValidationError(msg) + if "maxItems" in schema and len(value) > schema["maxItems"]: + msg = f"{path}: expected at most {schema['maxItems']} items, got {len(value)}" + raise ValidationError(msg) + + prefix = schema.get("prefixItems", ()) + items = schema.get("items") + result = [] + for i, item in enumerate(value): + item_schema = prefix[i] if i < len(prefix) else items + result.append( + validate(item, item_schema, f"{path}[{i}]") if item_schema else item + ) + return result + + +def _check_object(value, schema, path): + properties = schema.get("properties", {}) + additional = schema.get("additionalProperties", True) + + missing = [k for k in schema.get("required", ()) if k not in value] + if missing: + msg = f"{path}: missing required {', '.join(repr(k) for k in missing)}" + raise ValidationError(msg) + + result = {} + for key, item in value.items(): + if not isinstance(key, str): + msg = f"{path}: keys must be strings" + raise ValidationError(msg) + if key in properties: + result[key] = validate(item, properties[key], f"{path}.{key}") + elif additional is False: + msg = f"{path}: unexpected {key!r}" + ( + f" (allowed: {sorted(properties)})" if properties else "" + ) + raise ValidationError(msg) + elif isinstance(additional, dict): + result[key] = validate(item, additional, f"{path}.{key}") + else: + result[key] = item + return result + + +def bind_arguments(fn, arguments, schema, label="arguments", positional=()): + """Validate MCP call arguments for ``fn`` and return the coerced dict. + + Checks them against ``schema`` and against ``fn``'s signature (catching + missing/unexpected arguments the schema may not describe), so that a + ``TypeError`` raised later can only come from the function body. + ``positional`` are the leading arguments the caller supplies itself + (a constructor's server). Object-typed parameters keep the ids or + names given; see :func:`resolve_object_arguments`. + """ + if arguments is None: + arguments = {} + if not isinstance(arguments, dict): + msg = f"{label}: expected object, got {_type_name(arguments)}" + raise ValidationError(msg) + arguments = validate(arguments, schema, label) + # The Progress parameter is filled in later (see progress.inject) + reporter = progress_parameter(fn) + filled = {reporter: None} if reporter and reporter not in arguments else {} + try: + inspect.signature(fn).bind(*positional, **arguments, **filled) + except TypeError as e: + msg = f"{label}: {e}" + raise ValidationError(msg) from None + hints = parameter_hints(fn) + return { + name: from_json(value, hints.get(name)) for name, value in arguments.items() + } + + +def is_object(value) -> bool: + """Whether ``value`` is an object agents may reach, once its class declares + something: an instance of a class that inherits ``MCPExposable``, such as + a component or a trame-dataclass model.""" + return isinstance(value, MCPExposable) + + +def object_class(annotation): + """``annotation`` when it is a class whose instances agents may reach, + else None.""" + if isinstance(annotation, type) and issubclass(annotation, MCPExposable): + return annotation + return None + + +def object_classes(annotation) -> set: + """The component and model classes named anywhere in ``annotation``.""" + found, stack = set(), [annotation] + while stack: + current = stack.pop() + cls = object_class(current) + if cls is not None: + found.add(cls) + else: + stack.extend(typing.get_args(current)) + return found + + +def resolve_object_arguments(fn, arguments, resolve, label="arguments"): + """``arguments`` (bound by :func:`bind_arguments`) with the ids or names given + for object-typed parameters replaced by the objects: + ``resolve(id, cls, path)`` returns the object or raises.""" + hints = parameter_hints(fn) + typed = {name for name, hint in hints.items() if object_classes(hint)} + if not typed & arguments.keys(): + return arguments + return { + name: from_json(value, hints[name], resolve, f"{label}.{name}") + if name in typed + else value + for name, value in arguments.items() + } diff --git a/src/trame_server/mcp/typeinfo.py b/src/trame_server/mcp/typeinfo.py new file mode 100644 index 0000000..51faaf7 --- /dev/null +++ b/src/trame_server/mcp/typeinfo.py @@ -0,0 +1,448 @@ +"""What agents see of a class: its properties, methods and constructor. + +trame-common and trame-dataclass only store what a class declares, as plain +data (``Cls.mcp_describe()``); this module interprets it. A class is visible +when a keyword or ``mcp_describe()`` declares it, or when it has markers +(``@mcp.method``, ``@mcp.property``), an exposed field or a constructor, +unless it is hidden with False. Fields and constructors come from the +class's ``MCPExposable`` hooks (``_mcp_fields``, ``_mcp_constructor``), +which apply the rules of its type of class, such as a model's +``all_fields``. +""" + +from __future__ import annotations + +import inspect +import logging +import typing +import weakref +from dataclasses import dataclass, field + +from trame_common.obj import mcp as declarations +from trame_common.obj.mcp import METHOD_MARKER, PROPERTY_MARKER + +from trame_server.mcp.schema import ( + check_progress, + docstring, + nullable, + object_class, + object_classes, + return_hint, + type_to_schema, + unwrap, +) + +logger = logging.getLogger("trame_server.mcp") + +__all__ = [ + "Constructor", + "Method", + "Property", + "TypeInfo", + "is_visible", + "property_schema", + "reference_schema", + "tool_reference_schema", + "type_info", + "type_name", + "type_names", +] + +STATE_TYPE = "State" + + +# Methods trame calls back on its own (@change, @trigger, @controller, @watch) +_CALLBACK_MARKERS = ( + "_trame_state_change", + "_trame_trigger_names", + "_trame_controller", + "_watch", +) + + +@dataclass(frozen=True) +class Property: + name: str + origin: str # "field" (_mcp_fields), "property" (Python) or "state" (key) + schema: dict # what is written (objects by id), without its description + description: str | None = None + writable: bool = False # can be written, permissions aside + requires: tuple = () # to read it at all + set_requires: tuple = () # to write it, in addition + reference: bool = False # may hold objects agents reach + type_def: object = None # the declared type (tuples, objects by id) + + def described(self) -> dict: + """The schema with the description, as agents read it.""" + if not self.description: + return dict(self.schema) + return {**self.schema, "description": self.description} + + +@dataclass(frozen=True) +class Method: + name: str # the name agents call it by + attr: str + description: str + read_only: bool = False + annotations: dict | None = None # MCP tool annotations, but readOnlyHint + concurrent: bool = False + requires: tuple = () + + +@dataclass(frozen=True) +class Constructor: + attr: str | None # the classmethod, or None for cls(server, **values) + description: str = "" + requires: tuple = () + server: str | None = None # the classmethod's parameter taking the server + + +@dataclass(frozen=True) +class TypeInfo: + cls: type | None # None for the state object + visible: bool + description: str | None = None + properties: dict = field(default_factory=dict) + methods: dict = field(default_factory=dict) + constructor: Constructor | None = None + # Why the class's constructor cannot work, warned about when described + constructor_warning: str | None = None + # Reference fields left hidden, name -> classes they hold (startup audit) + hidden_references: dict = field(default_factory=dict) + + +HIDDEN = TypeInfo(cls=None, visible=False) + +# Cached on the class itself (its own __dict__, never inherited): a +# TypeInfo refers to its class, which a WeakKeyDictionary would keep alive +_CACHE_ATTR = "_trame_mcp_type_info" +_VISIBLE = weakref.WeakKeyDictionary() # cls -> (version, bool) + + +def type_info(cls) -> TypeInfo: + """What agents see of ``cls`` (cached until a class declares or is + described again).""" + version = declarations.options_version() + cached = vars(cls).get(_CACHE_ATTR) + if cached is not None and cached[0] == version: + return cached[1] + info = _build(cls) if is_visible(cls) else HIDDEN + setattr(cls, _CACHE_ATTR, (version, info)) + return info + + +def is_visible(cls) -> bool: + """Whether agents may see ``cls``, a component, model or other ``MCPExposable`` + class.""" + if not isinstance(cls, type): + return False + version = declarations.options_version() + cached = _VISIBLE.get(cls) + if cached is not None and cached[0] == version: + return cached[1] + visible = _declares(cls) + _VISIBLE[cls] = (version, visible) + return visible + + +def _declares(cls): + if object_class(cls) is None: + return False + declared = declarations.class_options(cls)["declared"] + if declared is not None: + return declared + # A field's mcp=True is an empty dict of options, which still exposes it + return ( + declarations.has_markers(cls) + or any(isinstance(f.options, dict) for f in cls._mcp_fields().values()) + or cls._mcp_constructor() is not None + ) + + +def _object_reference(annotation, describe): + """The schema of a value that holds an object of the class ``annotation``. It is + ``describe(name)`` for a visible class, and a plain string otherwise, since + agents cannot use the id. None when ``annotation`` names no such class.""" + if typing.get_origin(annotation) is not None: + return None # on Python 3.10, isinstance(list[X], type) holds + cls = object_class(annotation) + if cls is None: + return None + if not is_visible(cls): + return {"type": "string"} + return describe(type_name(cls)) + + +def reference_schema(annotation): + """The ``type_to_schema`` hook that describes a value holding a component or + model as ``{"object": "T"}``, the same in properties, parameters and returns. + Agents read it as an ``{"id", "type"}`` reference, a method returns it as + object_get shows it, and agents write it by id (a parameter also takes an + object name). ``validate`` checks that it is a string.""" + return _object_reference(annotation, lambda name: {"object": name}) + + +def tool_reference_schema(annotation): + """The ``type_to_schema`` hook for the input schema of an app's own tool, as + clients read it in ``tools/list``. An object parameter is a string, and its + description says what type of object it names.""" + return _object_reference( + annotation, + lambda name: {"type": "string", "description": f"{name} object, by id or name"}, + ) + + +# ----------------------------------------------------------------------------- +# Type names +# ----------------------------------------------------------------------------- + +_NAMES = {"version": None, "names": {}, "classes": {}} +_WARNED = weakref.WeakSet() + + +def type_names(): + """``(names, classes)``: the agent-facing name of every visible class, and + the reverse. A class is named after itself; when an earlier visible class + already has that name, it gets its module-qualified name, with a + warning. ``State`` is the state object's.""" + version = declarations.options_version() + if _NAMES["version"] == version: + return _NAMES["names"], _NAMES["classes"] + + names, classes = weakref.WeakKeyDictionary(), weakref.WeakValueDictionary() + for cls in declarations.declared_classes(): + if not is_visible(cls): + continue + name = cls.__name__ + if name in classes or name == STATE_TYPE: + first = classes.get(name) + name = f"{cls.__module__}.{cls.__qualname__}" + if cls not in _WARNED: + _WARNED.add(cls) + logger.warning( + "Two classes agents see are named %s: %s keeps the name, %s is " + "called %s", + cls.__name__, + "the state object" if first is None else first.__qualname__, + cls.__qualname__, + name, + ) + names[cls] = name + classes[name] = cls + _NAMES.update(version=version, names=names, classes=classes) + return names, classes + + +def type_name(cls) -> str: + return type_names()[0].get(cls) or cls.__name__ + + +# ----------------------------------------------------------------------------- +# Building +# ----------------------------------------------------------------------------- + + +def _members(cls): + """``{name: member}`` of ``cls`` and its bases, in definition order, a base's + members first. A member that a subclass overrides keeps the place where the + base defined it.""" + members = {} + for klass in reversed(cls.__mro__): + members.update(vars(klass)) + return members + + +def _base_names(cls): + """The names of the members that trame's own base classes define. Those are the + classes from the packages that the classes in the MRO name in + ``_mcp_api_packages``. ``all_methods`` never exposes these members, and + marking or listing one is an error.""" + packages = {"builtins"} + for klass in cls.__mro__: + packages.update(vars(klass).get("_mcp_api_packages", ())) + names = set() + for klass in cls.__mro__: + if klass.__module__.split(".")[0] in packages: + names.update(vars(klass)) + return names + + +def property_schema(annotation, options, *, ref=None, can_be_null=False): + """``(schema, description)`` of a property from its annotation and its + options. The options' ``schema`` keywords go over the annotation's + schema, which accepts null too with ``can_be_null``. The description is + the options' own, else the annotation's (``Annotated[T, "..."]``), and + never stays in the schema.""" + schema = type_to_schema(annotation, ref) + schema.update(options.get("schema") or {}) + found = schema.pop("description", None) + if can_be_null: + schema = nullable(schema) + return schema, options.get("description") or found + + +def _build(cls): + options = declarations.class_options(cls) + fields = cls._mcp_fields() + base = _base_names(cls) + properties, hidden_references = {}, {} + + # Fields (a model's), as the class's _mcp_fields hook exposes them + for name, f in fields.items(): + opts = f.options + if not isinstance(opts, dict): + if opts is None and f.holds_objects: + hidden_references[name] = object_classes(f.type_def) + continue + schema, description = property_schema( + f.type_def, opts, ref=reference_schema, can_be_null=f.nullable + ) + properties[name] = Property( + name=name, + origin="field", + schema=schema, + description=description, + writable=not opts.get("read_only", False), + requires=opts.get("requires", ()), + reference=f.holds_objects, + type_def=f.type_def, + ) + + # Python properties: marked, or listed + for attr, member in _members(cls).items(): + if attr in fields: + continue + if not isinstance(member, property): + continue + marker = getattr(member, PROPERTY_MARKER, None) + listed = options["properties"].get(attr) + if listed is False or (marker is None and listed is None): + continue + if attr in base: + msg = f"{cls.__name__}.{attr} is part of trame's API and cannot be exposed" + raise ValueError(msg) + opts = {**(marker or {}), **(listed or {})} + hint = return_hint(member.fget) if member.fget else None + schema, description = property_schema(hint, opts, ref=reference_schema) + properties[attr] = Property( + name=attr, + origin="property", + schema=schema, + description=description or docstring(member.fget), + writable=member.fset is not None and not opts.get("read_only"), + requires=opts.get("requires", ()), + set_requires=opts.get("set_requires", ()), + reference=bool(object_classes(hint)), + type_def=hint, + ) + + constructor, warning = _constructor(cls) + return TypeInfo( + cls=cls, + visible=True, + description=_description(cls, options), + properties=properties, + methods=_methods(cls, options, base), + constructor=constructor, + constructor_warning=warning, + hidden_references=hidden_references, + ) + + +def _description(cls, options): + """The class's own description option, else its own docstring; neither + is inherited.""" + if options["description"]: + return options["description"] + doc = vars(cls).get("__doc__") + return unwrap(inspect.cleandoc(doc)) if doc else None + + +def _methods(cls, options, base): + methods = {} + for attr, fn in _members(cls).items(): + if not inspect.isfunction(fn): + continue # classmethods, staticmethods, properties, data + marker = fn.__dict__.get(METHOD_MARKER) + listed = options["methods"].get(attr) + if listed is False: + continue + automatic = ( + options["all_methods"] + and not attr.startswith("_") + and attr not in base + and not any(m in fn.__dict__ for m in _CALLBACK_MARKERS) + ) + if marker is None and listed is None and not automatic: + continue + if attr in base: + msg = f"{cls.__name__}.{attr} is part of trame's API and cannot be exposed" + raise ValueError(msg) + check_progress(fn, f"{cls.__name__}.{attr}") + opts = {**(marker or {}), **(listed or {})} + public = opts.get("name") or attr + if public in methods: + msg = f"{cls.__name__}: two methods are exposed as {public!r}" + raise ValueError(msg) + annotations = dict(opts.get("annotations") or {}) + read_only = annotations.pop("readOnlyHint", False) or opts.get("read_only") + methods[public] = Method( + name=public, + attr=attr, + description=opts.get("description") or docstring(fn) or "", + read_only=bool(read_only), + annotations=annotations or None, + concurrent=bool(opts.get("concurrent", False)), + requires=tuple(opts.get("requires", ())), + ) + return methods + + +def _constructor(cls): + """``(Constructor or None, warning or None)``, from the class's + ``_mcp_constructor`` hook.""" + declared = cls._mcp_constructor() + if declared is None: + return None, None + if declared.problem: + return None, declared.problem + if declared.attr is None: + constructor = Constructor( + attr=None, + description=declared.description, + requires=tuple(declared.requires), + ) + return constructor, None + + attr = declared.attr + member = inspect.getattr_static(cls, attr, None) + if not isinstance(member, classmethod): + msg = f"{cls.__name__}.{attr} is the constructor but is not a classmethod" + raise ValueError(msg) + parameters = list(inspect.signature(getattr(cls, attr)).parameters) + if not parameters: + msg = ( + f"{cls.__name__}.{attr}: a constructor takes the trame server as its " + f"first parameter after cls" + ) + raise ValueError(msg) + return ( + Constructor( + attr=attr, + description=declared.description or docstring(member.__func__) or "", + requires=tuple(declared.requires), + server=parameters[0], + ), + None, + ) + + +def subclasses(cls): + """``cls`` and all its subclasses.""" + result, stack = [], [cls] + while stack: + klass = stack.pop() + result.append(klass) + stack.extend(klass.__subclasses__()) + return result diff --git a/src/trame_server/protocol.py b/src/trame_server/protocol.py index 4c3743f..dfe0e23 100644 --- a/src/trame_server/protocol.py +++ b/src/trame_server/protocol.py @@ -93,6 +93,7 @@ def initialize(self): # Called by wslink def set_server(self, _server): self.server._server = _server + self.server._bind_mcp(_server) if self.server.controller.on_server_bind.exists(): # wslink registers its static routes on start, so routes added # here take precedence over them @@ -114,6 +115,9 @@ def port_callback(self, port_used): if self.server.controller.on_server_ready.exists(): self.server.controller.on_server_ready(**self.server.state.to_dict()) + # Announce agent access (after the connection info printed above) + self.server._mcp_ready() + # Mark the server as ready if not self.server.ready.done(): self.server.ready.set_result(True) diff --git a/src/trame_server/state.py b/src/trame_server/state.py index a5ebd57..b5a1b2f 100644 --- a/src/trame_server/state.py +++ b/src/trame_server/state.py @@ -180,8 +180,12 @@ def __init__( commit_fn=None, hot_reload=False, ready=False, + get_mcp=None, ): self._push_state_fn = commit_fn + # Returns the server's MCP; None for a state agents never see + # (server.context) + self._get_mcp = get_mcp self._hot_reload = hot_reload self._translator = translator or Translator() self._modified_keys = share(internal, "_modified_keys", set()) @@ -516,6 +520,46 @@ def register_change_callback(func): return register_change_callback + # ------------------------------------------------------------------------- + # Agents (MCP) + # ------------------------------------------------------------------------- + + def _mcp(self): + if self._get_mcp is None: + msg = ( + "This state cannot be exposed to agents (server.context is " + "server-only); expose keys of server.state instead" + ) + raise RuntimeError(msg) + return self._get_mcp() + + def mcp_expose( + self, name, *, type=None, schema=None, description=None, read_only=False + ): + """ + Expose a state key to agents, as a property of the ``state`` object. + Writes are checked against ``type`` and ``schema`` and applied inside + ``with state:``, so ``@state.change`` reactions run and browsers update. + On a child server the key is translated: agents see the real key. + + :param name: the state key + :param type: Python type of its values, e.g. ``int`` or ``list[float]`` + :param schema: extra JSON Schema keywords (``minimum``, ``enum``, ...) + :param description: what agents read about it + :param read_only: agents may read it but not write it + """ + self._mcp()._expose_state( + self._translator.translate_key(name), + type=type, + schema=schema, + description=description, + read_only=read_only, + ) + + def mcp_unexpose(self, *names): + """Withdraw state keys from agents.""" + self._mcp()._unexpose_state(*self._translator.translate_list(names)) + @contextmanager def suppress_change_listeners(self, *keys: str): """ diff --git a/src/trame_server/utils/server.py b/src/trame_server/utils/server.py index ae2ef4c..991a0b9 100644 --- a/src/trame_server/utils/server.py +++ b/src/trame_server/utils/server.py @@ -25,3 +25,26 @@ def print_informations(server): flush=True, ) sys.stdout.flush() + + +def print_mcp_informations(server): + """Print how agents connect: the MCP URL, and the commands that start a Claude + Code or Codex session with this server added.""" + from trame_server.mcp import claude_command, codex_command # noqa: PLC0415 + + mcp = server.root_server.mcp + http = mcp.http + url, name = http.url, mcp.name + # A generated token is shown in full, since nobody else knows it. A token + # the user set is shown as a placeholder + token = http.token if http._token_generated else ("" if http.token else None) + print() + print(f"MCP server available at {url}") + print(f" {claude_command(name, url, token)}") + print(f" {codex_command(name, url, token)}") + if http._token_generated: + print( + " (token generated because the server listens beyond localhost; " + "set it with --mcp-http-token or TRAME_MCP_HTTP_TOKEN)" + ) + print(flush=True) diff --git a/tests/test_mcp_call_log.py b/tests/test_mcp_call_log.py new file mode 100644 index 0000000..61988eb --- /dev/null +++ b/tests/test_mcp_call_log.py @@ -0,0 +1,300 @@ +"""The call log (``--mcp-log``): what agents call, logged once whichever way +the call came in, through the standard ``logging`` module.""" + +import asyncio +import itertools +import logging +import re + +import pytest +from trame.app import get_server + +from trame_server.mcp import Image +from trame_server.mcp.call_log import shorten + +SESSION = "abcdef0123456789" +LAUNCH_VARIABLES = ( + "TRAME_ARGS", + "TRAME_MCP_HTTP_TOKEN", + "TRAME_MCP_LOG", + "TRAME_MCP_DENY", +) +MCP_LOGGER = logging.getLogger("trame_server.mcp") + +_counter = itertools.count() + + +def rpc(method, params=None, rpc_id=1): + msg = {"jsonrpc": "2.0", "id": rpc_id, "method": method} + if params is not None: + msg["params"] = params + return msg + + +class _Lines(logging.Handler): + def __init__(self): + super().__init__() + self.lines = [] + + def emit(self, record): + self.lines.append(record.getMessage()) + + +@pytest.fixture(autouse=True) +def _no_launch_options(monkeypatch): + for name in LAUNCH_VARIABLES: + monkeypatch.delenv(name, raising=False) + + +@pytest.fixture +def pristine_logger(): + """The MCP loggers as an application that never configured logging has + them (WARNING, Python's default; pytest lowers the root logger's), + restored afterwards (they are process-wide).""" + saved = MCP_LOGGER.handlers[:], MCP_LOGGER.level + MCP_LOGGER.handlers.clear() + MCP_LOGGER.setLevel(logging.WARNING) + yield MCP_LOGGER + MCP_LOGGER.handlers[:], MCP_LOGGER.level = saved + + +@pytest.fixture +def lines(pristine_logger): + """Every line the MCP loggers write, through a handler of the + application's own.""" + handler = _Lines() + pristine_logger.addHandler(handler) + return handler.lines + + +@pytest.fixture +def logging_on(pristine_logger): + """What --mcp-log does: the MCP loggers at INFO.""" + pristine_logger.setLevel(logging.INFO) + + +def new_server(): + return get_server(f"test_mcp_call_log_{next(_counter)}") + + +@pytest.fixture +def mcp(pristine_logger): # noqa: ARG001 - whatever a test turns on is undone + mcp = new_server().mcp + mcp.enabled = True + return mcp + + +@pytest.fixture +def tools(mcp): + @mcp.tool + def add(a: int, b: int = 1) -> int: + return a + b + + @mcp.tool + def failing(): + msg = "boom" + raise RuntimeError(msg) + + @mcp.tool + def echo(text: str) -> str: + return text + + @mcp.tool + def snapshot() -> list: + return [Image(b"\x89PNG" * 10), "a view"] + + @mcp.resource("trame://info", mime_type="application/json") + def info(): + return {"name": "test"} + + @mcp.prompt + def walkthrough(topic: str = "app") -> str: + return f"Explain the {topic}" + + return mcp + + +@pytest.mark.asyncio +async def test_off_by_default(tools, lines): + await tools.protocol.handle( + rpc("tools/call", {"name": "add", "arguments": {"a": 2}}) + ) + assert lines == [] + + +@pytest.mark.asyncio +async def test_a_call_is_logged_once_when_it_ends(tools, lines, logging_on): # noqa: ARG001 + response = await tools.protocol.handle( + rpc("tools/call", {"name": "add", "arguments": {"a": 2}}), session=SESSION + ) + assert response["result"]["content"][0]["text"] == "3" + (line,) = lines + assert re.fullmatch(r"mcp\[abcdef\] add \{a: 2\} -> ok \d+\.\d\ds", line), line + + +@pytest.mark.asyncio +async def test_resources_and_prompts_too(tools, lines, logging_on): # noqa: ARG001 + await tools.protocol.handle(rpc("resources/read", {"uri": "trame://info"})) + await tools.protocol.handle(rpc("prompts/get", {"name": "walkthrough"})) + resource, prompt = lines + assert re.fullmatch(r"mcp\[-\] resource trame://info -> ok \d+\.\d\ds", resource) + assert re.fullmatch(r"mcp\[-\] prompt walkthrough -> ok \d+\.\d\ds", prompt) + + +@pytest.mark.asyncio +async def test_in_process_calls_are_logged_too(tools, lines, logging_on): # noqa: ARG001 + client = tools.client() + await client.call_tool("add", {"a": 1}) + await client.read_resource("trame://info") + assert lines[0].startswith("mcp[in-process] add {a: 1} -> ok") + assert lines[1].startswith("mcp[in-process] resource trame://info -> ok") + + +@pytest.mark.asyncio +async def test_clients_without_a_session_are_named_by_themselves( + tools, + lines, + logging_on, # noqa: ARG001 +): + """In-process clients and clients of MCP 2026-07-28 have no session id: + the name they give (clientInfo) says who called, kept to safe + characters.""" + await tools.client().call_tool("add", {"a": 1}) + await tools.client("assistant").call_tool("add", {"a": 1}) + await tools.client("my agent\n[x]").call_tool("add", {"a": 1}) + meta = { + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientCapabilities": {}, + "io.modelcontextprotocol/clientInfo": {"name": "claude-code", "version": "2"}, + } + await tools.protocol.handle( + rpc("tools/call", {"name": "add", "arguments": {"a": 1}, "_meta": meta}) + ) + callers = [line.split("]")[0] + "]" for line in lines] + assert callers == [ + "mcp[in-process]", + "mcp[assistant]", + "mcp[my_agent__x_]", # no newline or bracket of the name's own + "mcp[claude-code]", + ] + + +@pytest.mark.asyncio +async def test_failures_say_what_the_agent_saw(tools, lines, logging_on): # noqa: ARG001 + await tools.protocol.handle(rpc("tools/call", {"name": "failing"})) + await tools.protocol.handle( + rpc("tools/call", {"name": "add", "arguments": {"a": "x"}}) + ) + await tools.protocol.handle( + rpc("tools/call", {"name": "missing", "arguments": {"q": 1}}) + ) + await tools.protocol.handle(rpc("resources/read", {"uri": "trame://nope"})) + # One line each: the failure is not logged a second time + failing, invalid, unknown, resource = ( + line for line in lines if not line.startswith("Tool failing raised") + ) + assert re.match( + r"mcp\[-\] failing -> error \d+\.\d\ds: RuntimeError: boom$", failing + ) + assert 'add {a: "x"} -> error' in invalid + assert invalid.endswith( + "Invalid arguments for tool add: arguments.a: expected integer, got string" + ) + assert unknown.startswith("mcp[-] missing {q: 1} -> error") + assert unknown.endswith("Unknown tool: missing") + assert resource.startswith("mcp[-] resource trame://nope -> error") + assert resource.endswith("Unknown resource: trame://nope") + + +@pytest.mark.asyncio +async def test_a_refused_permission_is_logged_as_denied(tools, lines, logging_on): # noqa: ARG001 + tools.permissions.declare("files", "Read files", allowed=False) + + @tools.tool(requires="files") + def read(path: str) -> str: + return path + + await tools.protocol.handle( + rpc("tools/call", {"name": "read", "arguments": {"path": "/"}}) + ) + (line,) = lines + assert re.match( + r'mcp\[-\] read \{path: "/"\} -> denied \d+\.\d\ds: read: permission "files" ' + r"\(Read files\) is off", + line, + ), line + + +@pytest.mark.asyncio +async def test_disabled_access_is_logged(tools, lines, logging_on): # noqa: ARG001 + tools.enabled = False + await tools.protocol.handle( + rpc("tools/call", {"name": "add", "arguments": {"a": 1}}) + ) + (line,) = lines + assert "-> error" in line + assert line.endswith("MCP access is currently disabled by the application") + + +@pytest.mark.asyncio +async def test_a_cancelled_call(mcp, lines, logging_on): # noqa: ARG001 + started = asyncio.Event() + + @mcp.tool + async def wait_forever(): + started.set() + await asyncio.Event().wait() + + handling = asyncio.ensure_future( + mcp.protocol.handle(rpc("tools/call", {"name": "wait_forever"}, rpc_id=7)) + ) + await started.wait() + await mcp.protocol.handle( + { + "jsonrpc": "2.0", + "method": "notifications/cancelled", + "params": {"requestId": 7}, + } + ) + await handling + calls = [line for line in lines if line.startswith("mcp[")] + (line,) = calls + assert re.fullmatch(r"mcp\[-\] wait_forever -> cancelled \d+\.\d\ds", line), line + + +@pytest.mark.asyncio +async def test_long_arguments_are_shortened(tools, lines, logging_on): # noqa: ARG001 + await tools.protocol.handle( + rpc("tools/call", {"name": "echo", "arguments": {"text": "x" * 500}}) + ) + assert lines[0].startswith("mcp[-] echo {text: } -> ok") + + +def test_shorten(): + assert ( + shorten({"a": 1, "b": [True, None], "c": "hi"}) + == '{a: 1, b: [true, null], c: "hi"}' + ) + assert shorten(list(range(10))) == "[0, 1, 2, 3, 4, 5, 6, 7, +2]" + assert shorten({str(i): i for i in range(9)}).endswith(", 7: 7, +1}") + assert shorten({"a": {"b": {"c": {"d": 1}}}}) == "{a: {b: {c: {...}}}}" + assert shorten([[[[1]]]]) == "[[[[...]]]]" + assert shorten(object()) == "" + assert shorten(1.5) == "1.5" + + +@pytest.mark.parametrize("variable", ["TRAME_ARGS", "TRAME_MCP_LOG"]) +def test_the_launch_flag_sets_the_level(monkeypatch, pristine_logger, variable): + # trame reads TRAME_ARGS instead of sys.argv when running under pytest + monkeypatch.setenv(variable, "--mcp-log" if variable == "TRAME_ARGS" else "1") + new_server().mcp # noqa: B018 - created, it applies the launch options + assert pristine_logger.level == logging.INFO + # A level only: where the lines go is the application's choice + assert pristine_logger.handlers == [] + + +def test_the_launch_flag_keeps_a_finer_level(monkeypatch, pristine_logger): + pristine_logger.setLevel(logging.DEBUG) + monkeypatch.setenv("TRAME_ARGS", "--mcp-log") + new_server().mcp # noqa: B018 + assert pristine_logger.level == logging.DEBUG diff --git a/tests/test_mcp_client.py b/tests/test_mcp_client.py new file mode 100644 index 0000000..7d27a56 --- /dev/null +++ b/tests/test_mcp_client.py @@ -0,0 +1,209 @@ +"""Agents inside the application's process (``server.mcp.client()``): the +same rules as agents connected over a transport, with no transport.""" + +import asyncio +import itertools + +import pytest +from trame.app import get_server + +from trame_server.mcp import MCPClient, MCPError +from trame_server.mcp.protocol import INVALID_PARAMS, INVALID_REQUEST + +LAUNCH_VARIABLES = ( + "TRAME_ARGS", + "TRAME_MCP_HTTP_TOKEN", + "TRAME_MCP_LOG", + "TRAME_MCP_DENY", +) + +_counter = itertools.count() + +DISABLED = "MCP access is currently disabled by the application" + + +def new_server(): + return get_server(f"test_mcp_client_{next(_counter)}") + + +@pytest.fixture(autouse=True) +def _no_launch_options(monkeypatch): + for name in LAUNCH_VARIABLES: + monkeypatch.delenv(name, raising=False) + + +def declare(mcp): + @mcp.tool + def add(a: int, b: int = 1) -> int: + """Add two numbers.""" + return a + b + + @mcp.tool + def failing(): + msg = "boom" + raise RuntimeError(msg) + + @mcp.resource("trame://info", mime_type="text/plain") + def info(): + return "An app." + + @mcp.prompt + def review(topic: str) -> str: + """Review something.""" + return f"Review the {topic}" + + +@pytest.fixture +def mcp(): + mcp = new_server().mcp + mcp.configure(enabled=True, transports=[]) # no transport needed + declare(mcp) + return mcp + + +@pytest.mark.asyncio +async def test_tools(mcp): + client = mcp.client() + assert isinstance(client, MCPClient) + assert client.name == "in-process" + + tools = {tool["name"]: tool for tool in await client.list_tools()} + assert tools["add"]["description"] == "Add two numbers." + assert tools["add"]["inputSchema"]["required"] == ["a"] + + # The result as MCP defines it, without what only transports need + assert await client.call_tool("add", {"a": 2}) == { + "content": [{"type": "text", "text": "3"}] + } + + # A failing tool still answers: its result says it failed + failed = await client.call_tool("failing") + assert failed["isError"] is True + assert "boom" in failed["content"][0]["text"] + invalid = await client.call_tool("add", {"a": "x"}) + assert invalid["isError"] is True + assert "expected integer" in invalid["content"][0]["text"] + + +@pytest.mark.asyncio +async def test_what_the_server_cannot_serve_raises(mcp): + client = mcp.client() + with pytest.raises(MCPError, match="Unknown tool: missing") as caught: + await client.call_tool("missing") + assert caught.value.code == INVALID_PARAMS + with pytest.raises(MCPError, match="Unknown resource"): + await client.read_resource("trame://nope") + with pytest.raises(MCPError, match="Unknown prompt"): + await client.get_prompt("nope") + + +@pytest.mark.asyncio +async def test_resources_and_prompts(mcp): + client = mcp.client() + assert [r["uri"] for r in await client.list_resources()] == ["trame://info"] + assert await client.read_resource("trame://info") == [ + {"uri": "trame://info", "mimeType": "text/plain", "text": "An app."} + ] + + (prompt,) = await client.list_prompts() + assert prompt["name"] == "review" + assert await client.get_prompt("review", {"topic": "cone"}) == { + "description": "Review something.", + "messages": [ + {"role": "user", "content": {"type": "text", "text": "Review the cone"}} + ], + } + + +@pytest.mark.asyncio +async def test_instructions(mcp): + mcp.instructions = "A test app." + assert mcp.client().instructions == mcp._agent_instructions() + assert mcp.client().instructions.startswith("A test app.") + + +@pytest.mark.asyncio +async def test_access_follows_enabled(mcp): + client = mcp.client() + mcp.enabled = False + refused = await client.call_tool("add", {"a": 1}) + assert refused == {"content": [{"type": "text", "text": DISABLED}], "isError": True} + with pytest.raises(MCPError, match=DISABLED) as caught: + await client.read_resource("trame://info") + assert caught.value.code == INVALID_REQUEST + + mcp.enabled = True + assert (await client.call_tool("add", {"a": 1}))["content"][0]["text"] == "2" + + +@pytest.mark.asyncio +async def test_no_mcp_refuses_agents_in_the_process(monkeypatch): + # trame reads TRAME_ARGS instead of sys.argv when running under pytest + monkeypatch.setenv("TRAME_ARGS", "--no-mcp") + mcp = new_server().mcp + declare(mcp) + mcp.enabled = True # --no-mcp holds for the run + assert (await mcp.client().call_tool("add", {"a": 1}))["isError"] is True + + +@pytest.mark.asyncio +async def test_permissions_apply(mcp): + mcp.permissions.declare("math", "Do arithmetic") + + @mcp.tool(requires="math") + def double(x: int) -> int: + return 2 * x + + client = mcp.client() + assert (await client.call_tool("double", {"x": 2}))["content"][0]["text"] == "4" + mcp.permissions["math"] = False + assert "double" not in [t["name"] for t in await client.list_tools()] + assert (await client.call_tool("double", {"x": 2}))["isError"] is True + + +@pytest.mark.asyncio +async def test_cancelling_the_awaiting_task_cancels_the_call(mcp): + started, cleaned_up = asyncio.Event(), asyncio.Event() + + @mcp.tool + async def wait_forever() -> str: + started.set() + try: + await asyncio.Event().wait() + finally: + cleaned_up.set() + + call = asyncio.ensure_future(mcp.client().call_tool("wait_forever")) + await asyncio.wait_for(started.wait(), 1) + call.cancel() + with pytest.raises(asyncio.CancelledError): + await call + await asyncio.wait_for(cleaned_up.wait(), 1) + # the tool lock was released + assert (await mcp.client().call_tool("add", {"a": 1}))["content"][0]["text"] == "2" + + +@pytest.mark.asyncio +async def test_each_client_has_its_own_requests(mcp): + """Both clients number their requests from 1: cancelling one's request + leaves the other's alone.""" + mcp.concurrent_calls = True + started = {"a": asyncio.Event(), "b": asyncio.Event()} + release = asyncio.Event() + + @mcp.tool + async def wait(who: str) -> str: + started[who].set() + await release.wait() + return who + + a, b = mcp.client("a"), mcp.client("b") + call_a = asyncio.ensure_future(a.call_tool("wait", {"who": "a"})) + call_b = asyncio.ensure_future(b.call_tool("wait", {"who": "b"})) + await asyncio.wait_for(asyncio.gather(*(e.wait() for e in started.values())), 1) + + assert mcp.protocol.cancel(1, session=a) is True + release.set() + with pytest.raises(MCPError, match="Request cancelled"): + await call_a + assert (await call_b)["content"][0]["text"] == "b" diff --git a/tests/test_mcp_core.py b/tests/test_mcp_core.py new file mode 100644 index 0000000..ca6e77a --- /dev/null +++ b/tests/test_mcp_core.py @@ -0,0 +1,696 @@ +"""``server.mcp``: creation, options, first-class tools, resources and +prompts, and the launch options in ``Server.cli``.""" + +import itertools +import logging +from typing import Annotated + +import aiohttp +import pytest +from trame.app import get_server + +from trame_server.mcp import MCP, MCPDisabledError +from trame_server.mcp.objects import OBJECT_TOOLS + +LAUNCH_VARIABLES = ( + "TRAME_ARGS", + "TRAME_MCP_HTTP_TOKEN", + "TRAME_MCP_LOG", + "TRAME_MCP_DENY", +) + +_counter = itertools.count() + + +def new_server(): + return get_server(f"test_mcp_core_{next(_counter)}") + + +@pytest.fixture(autouse=True) +def _no_launch_options(monkeypatch): + for name in LAUNCH_VARIABLES: + monkeypatch.delenv(name, raising=False) + + +@pytest.fixture +def server(): + return new_server() + + +@pytest.fixture +def mcp(server): + mcp = server.mcp + mcp.enabled = True + return mcp + + +async def start(server, **kwargs): + server.start( + exec_mode="task", + port=0, + open_browser=False, + show_connection_info=False, + disable_logging=True, + **kwargs, + ) + await server.ready + + +# ----------------------------------------------------------------------------- +# Creation +# ----------------------------------------------------------------------------- + + +def test_created_on_first_access_one_per_root_server(server): + assert server._mcp is None + mcp = server.mcp + assert isinstance(mcp, MCP) + assert server.mcp is mcp + assert mcp.server is server + assert mcp.protocol is mcp.protocol + + child = server.create_child_server(prefix="child_") + assert child.mcp is child.mcp # its root's, through its translator + assert child.mcp.protocol is mcp.protocol + assert child._mcp is None + + +def test_a_child_server_creates_its_root_s(): + root = new_server() + child = root.create_child_server(prefix="child_") + mcp = child.mcp + assert root._mcp is not None + assert mcp.server is root + assert mcp.name == root.name + + # Options set through it are the root's + mcp.max_response_tokens = 5000 + assert root.mcp.max_response_tokens == 5000 + + +def test_exposing_state_creates_it(server): + child = server.create_child_server(prefix="child_") + child.state.size = 3 + assert server._mcp is None + child.state.mcp_expose("size", type=int) + assert server._mcp is not None + assert server.mcp._exposes_anything() + child.state.mcp_unexpose("size") + assert not server.mcp._exposes_anything() + + +def test_the_context_is_never_exposed(server): + with pytest.raises(RuntimeError, match="server-only"): + server.context.mcp_expose("secret") + assert server._mcp is None + + +def test_off_by_default(server): + mcp = server.mcp + + @mcp.tool + def ping() -> str: + """Ping.""" + return "pong" + + server.state.x = 1 + server.state.mcp_expose("x") + assert mcp.enabled is False + assert mcp.enabled_locked is False + assert issubclass(MCPDisabledError, PermissionError) + + +def test_exposes_anything(server): + mcp = server.mcp + assert not mcp._exposes_anything() + mcp.resource("app://guide")(lambda: "guide") + assert mcp._exposes_anything() + + mcp = new_server().mcp + mcp.prompt(lambda: "hi", name="hello") + assert mcp._exposes_anything() + + mcp = new_server().mcp + mcp.tool(lambda: "hi", name="hello") + assert mcp._exposes_anything() + + +# ----------------------------------------------------------------------------- +# Options +# ----------------------------------------------------------------------------- + + +def test_defaults(server): + mcp = server.mcp + assert mcp.name == server.name + assert mcp.enabled is False + assert mcp.transports == ("http",) + assert mcp.http.path == "/mcp" + assert mcp.http.token is None + assert mcp.http.allowed_origins == () + assert mcp.concurrent_calls is False + assert mcp.structured_content is False + assert mcp.instructions is None + assert mcp.max_response_tokens == 8000 + assert dict(mcp.permissions) == {} + + +def test_configure_sets_options_and_none_leaves_them(server): + mcp = server.mcp + mcp.configure( + name="viewer", + enabled=True, + transports=[], + concurrent_calls=True, + structured_content=True, + instructions="A viewer.", + max_response_tokens=1000, + ) + assert mcp.name == "viewer" + assert mcp.enabled is True + assert mcp.transports == () + assert mcp.concurrent_calls is True + assert mcp.structured_content is True + assert mcp.instructions == "A viewer." + assert mcp.max_response_tokens == 1000 + + mcp.http.configure( + path="api/mcp", + token="s3cret", + allowed_origins=["https://a.example/", "https://b.example"], + ) + assert mcp.http.path == "/api/mcp" + assert mcp.http.token == "s3cret" + assert mcp.http.allowed_origins == ("https://a.example", "https://b.example") + + # None leaves options unchanged, False is a value + mcp.configure(name=None, enabled=None) + mcp.http.configure(token=None) + assert (mcp.name, mcp.enabled, mcp.http.token) == ("viewer", True, "s3cret") + mcp.configure(enabled=False) + mcp.http.configure(token="") + assert (mcp.enabled, mcp.http.token) == (False, None) + + with pytest.raises(TypeError): + mcp.configure(bogus=1) + with pytest.raises(TypeError): + mcp.configure(token="s3cret") # the HTTP transport's: mcp.http + + mcp.name = None # back to the server name + assert mcp.name == server.name + + +def test_max_response_tokens(server): + mcp = server.mcp + mcp.max_response_tokens = 500 + assert mcp.max_response_tokens == 500 + for bad in (499, 0, -1, "1000", 1000.0, True, None): + with pytest.raises( + ValueError, match="max_response_tokens must be an integer of at least 500" + ): + mcp.max_response_tokens = bad + with pytest.raises(ValueError, match="at least 500"): + mcp.configure(max_response_tokens=100) + assert mcp.max_response_tokens == 500 + + +def test_allowed_origins_take_one_or_several(server): + http = server.mcp.http + http.allowed_origins = "https://a.example/" + assert http.allowed_origins == ("https://a.example",) + http.allowed_origins = [] + assert http.allowed_origins == () + + +def test_transports(server): + mcp = server.mcp + mcp.transports = [] # agents in the process only + assert mcp.transports == () + mcp.transports = "http" + assert mcp.transports == ("http",) + mcp.transports = ["http", "http"] + assert mcp.transports == ("http",) + for bad in (["stdio"], ["none"], ["http", "HTTP"]): + with pytest.raises( + ValueError, match=r"Unknown MCP transport .*available: http" + ): + mcp.transports = bad + assert mcp.transports == ("http",) + + mcp._start() # it can change while the app runs, as enabled can + mcp.transports = [] + assert mcp.transports == () + assert mcp.transports_locked is False + + +# ----------------------------------------------------------------------------- +# Tools, resources and prompts +# ----------------------------------------------------------------------------- + + +def test_tool_decorator_bare_and_options(mcp): + @mcp.tool + def alpha(x: int) -> int: + """Add one.""" + return x + 1 + + @mcp.tool(name="beta_renamed", description="custom") + def beta() -> str: + return "ok" + + assert alpha(1) == 2 # the decorator returns the function + assert [t.name for t in mcp.list_tools()] == ["alpha", "beta_renamed"] + assert mcp.get_tool("alpha").description == "Add one." + assert mcp.get_tool("beta_renamed").description == "custom" + assert mcp.get_tool("beta") is None + schema = mcp._input_schema(mcp.get_tool("alpha")) + assert schema["required"] == ["x"] + assert schema["properties"] == {"x": {"type": "integer"}} + + +def test_bound_methods_are_tools_resources_and_prompts(mcp): + class CoreViewer: + def screenshot(self, width: int = 300) -> str: + """Capture the view.""" + return f"{width}px" + + def guide(self) -> str: + """How to use the viewer.""" + return "# Guide" + + def review_session(self, focus: Annotated[str, "What to review"]) -> str: + return f"Review {focus}" + + viewer = CoreViewer() + mcp.tool(viewer.screenshot) + mcp.resource("viewer://guide", mime_type="text/markdown")(viewer.guide) + mcp.prompt(viewer.review_session) + + tool = mcp.get_tool("screenshot") + assert tool.description == "Capture the view." + assert tool.fn() == "300px" + assert list(mcp._input_schema(tool)["properties"]) == ["width"] + + resource = mcp.get_resource("viewer://guide") + assert (resource.name, resource.mime_type) == ("guide", "text/markdown") + assert resource.description == "How to use the viewer." + assert mcp.list_resources() == [resource] + + prompt = mcp.get_prompt("review_session") + assert prompt.arguments == [ + {"name": "focus", "required": True, "description": "What to review"} + ] + assert mcp.list_prompts() == [prompt] + + +def test_prompt_keeps_its_description_apart_from_its_arguments(mcp): + @mcp.prompt(description="Explore a file.") + def explore(path: Annotated[str, "File to explore"], depth: int = 1) -> str: + return path * depth + + @mcp.prompt + def documented(path: Annotated[str, "File to explore"], *args, **kwargs) -> str: # noqa: ARG001 + """Explore with the docstring.""" + return path + + explore_entry = mcp.get_prompt("explore") + assert explore_entry.description == "Explore a file." + assert explore_entry.arguments == [ + {"name": "path", "required": True, "description": "File to explore"}, + {"name": "depth", "required": False}, + ] + documented_entry = mcp.get_prompt("documented") + assert documented_entry.description == "Explore with the docstring." + assert [a["name"] for a in documented_entry.arguments] == ["path"] + + +def test_tool_annotations(mcp): + @mcp.tool(read_only=True) + def look() -> str: + return "ok" + + @mcp.tool(annotations={"destructiveHint": True}) + def wipe() -> str: + return "ok" + + @mcp.tool(read_only=True, annotations={"openWorldHint": False}) + def peek() -> str: + return "ok" + + @mcp.tool + def plain() -> str: + return "ok" + + assert mcp.get_tool("look").annotations == {"readOnlyHint": True} + assert mcp.get_tool("wipe").annotations == {"destructiveHint": True} + assert mcp.get_tool("peek").annotations == { + "openWorldHint": False, + "readOnlyHint": True, + } + assert mcp.get_tool("plain").annotations is None + assert mcp.get_tool("plain").concurrent is False + assert mcp.get_tool("plain").requires == () + + +def test_duplicate_names_raise(mcp): + @mcp.tool + def gamma(): + pass + + with pytest.raises(ValueError, match="A tool named 'gamma' is already registered"): + + @mcp.tool(name="gamma") + def other(): + pass + + def info(): + return {} + + mcp.resource("app://info")(info) + with pytest.raises(ValueError, match="'app://info' is already registered"): + mcp.resource("app://info")(info) + + mcp.prompt(lambda: "", name="intro") + with pytest.raises( + ValueError, match="A prompt named 'intro' is already registered" + ): + mcp.prompt(lambda: "", name="intro") + + +RESERVED = ( + "object_list", + "object_types", + "object_get", + "object_update", + "object_call", + "object_batch", +) + + +def test_the_reserved_names_are_the_object_tools(): + assert set(OBJECT_TOOLS) == set(RESERVED) + + +@pytest.mark.parametrize("name", RESERVED) +def test_the_object_tool_names_are_reserved(mcp, name): + def fn(): + pass + + with pytest.raises(ValueError, match="is the name of a built-in object tool"): + mcp.tool(fn, name=name) + fn.__name__ = name + with pytest.raises(ValueError, match="is the name of a built-in object tool"): + mcp.tool(fn) + + # also once the object tools are there + mcp.server.state.size = 1 + mcp.server.state.mcp_expose("size") + entry = mcp.get_tool(name) + with pytest.raises(ValueError, match="is the name of a built-in object tool"): + mcp.tool(fn) + assert mcp.get_tool(name) is entry + + +def test_tool_names_use_the_characters_clients_accept(mcp, caplog): + def fn(): + pass + + for bad in ("has space", "émoji", "a/b", 3): + with pytest.raises(ValueError, match="A tool name is 1-128 of"): + mcp.tool(fn, name=bad) + mcp.tool(fn, name="with-dash_and_underscore") + + with caplog.at_level(logging.WARNING, logger="trame_server.mcp"): + mcp.tool(fn, name="with.dot") + assert "OpenAI-style function calling rejects dots in names" in caplog.text + assert mcp.get_tool("with.dot") is not None + + +def test_a_full_tool_name_over_64_characters_raises(mcp): + def fn(): + pass + + mcp.name = "app" # mcp__app__: 10 characters before the tool name + mcp.tool(fn, name="a" * 54) + with pytest.raises(ValueError, match="has 65 characters") as error: + mcp.tool(fn, name="b" * 55) + assert str(error.value) == ( + f"Tool '{'b' * 55}': its full name 'mcp__app__{'b' * 55}' has 65 characters; " + "clients such as Claude Code reject more than 64. Shorten the tool name " + "or mcp.name" + ) + assert mcp.get_tool("b" * 55) is None + + # the default name is the trame server's, which counts too + mcp = new_server().mcp + room = 64 - len(f"mcp__{mcp.name}__") + mcp.tool(fn, name="c" * room) + with pytest.raises(ValueError, match="has 65 characters"): + mcp.tool(fn, name="d" * (room + 1)) + + +def test_a_longer_name_set_later_is_checked_at_start(mcp, caplog): + @mcp.tool + def reset(): + pass + + mcp.name = "n" * 60 + with pytest.raises(ValueError, match=r"Shorten the tool name or mcp\.name"): + mcp.tool(lambda: None, name="other") + + with caplog.at_level(logging.WARNING, logger="trame_server.mcp"): + mcp._start() + assert ( + f"Tool 'reset': its full name 'mcp__{'n' * 60}__reset' has 72 characters" + in caplog.text + ) + assert mcp.get_tool("reset") is not None # warned, still served + + +# ----------------------------------------------------------------------------- +# Launch options (Server.cli) +# ----------------------------------------------------------------------------- + + +def test_the_launch_flags_are_part_of_the_server_cli(server): + cli = server.cli + args, unknown = cli.parse_known_args( + [ + "--mcp", + "--mcp-transport", + "http", + "--mcp-http-token", + "t0k", + "--mcp-log", + "--mcp-deny", + "files,net", + "--mcp-deny", + "scripts", + "--app-option", + ] + ) + assert args.mcp is True + assert args.mcp_transport == ("http",) + assert args.mcp_http_token == "t0k" + assert args.mcp_log is True + assert args.mcp_deny == ["files,net", "scripts"] + assert unknown == ["--app-option"] # harmless to apps that don't use MCP + + args, _ = cli.parse_known_args([]) + assert ( + args.mcp, + args.mcp_transport, + args.mcp_http_token, + args.mcp_log, + args.mcp_deny, + ) == (None, None, None, False, None) + assert cli.parse_known_args(["--no-mcp"])[0].mcp is False + + +def test_the_transport_flag(server, capsys): + cli = server.cli + + def transports(*argv): + return cli.parse_known_args(list(argv))[0].mcp_transport + + # comma-separated and repeatable, like --mcp-deny + assert transports("--mcp-transport", "none") == ("none",) + assert transports("--mcp-transport", "http,http") == ("http",) + assert transports("--mcp-transport", "http", "--mcp-transport", "http") == ("http",) + assert transports("--mcp-transport=http") == ("http",) + assert transports("--no-mcp-transport", "--mcp-transport", "none") == ("none",) + + assert cli.parse_known_args(["--no-mcp-transport"])[0].mcp_no_transport is True + assert cli.parse_known_args([])[0].mcp_no_transport is False + + # wrong values fail as they are parsed, like any other wrong value + for argv, error in [ + ( + ["--no-mcp-transport", "--mcp-transport", "http"], + "--no-mcp-transport serves none", + ), + ( + ["--mcp-transport", "http", "--no-mcp-transport"], + "--mcp-transport names a transport", + ), + (["--mcp-transport", "stdoi"], "unknown transport 'stdoi' (choose from http"), + (["--mcp-transport", "none,http"], "none can't be combined with a transport"), + (["--mcp-transport", "http", "--mcp-transport", "none"], "can't be combined"), + ]: + with pytest.raises(SystemExit): + cli.parse_known_args(argv) + assert error in capsys.readouterr().err + + +def test_launch_options_win_until_start(monkeypatch): + # trame reads TRAME_ARGS instead of sys.argv when running under pytest + monkeypatch.setenv("TRAME_ARGS", "--no-mcp --mcp-http-token from-cli") + mcp = new_server().mcp + mcp.configure(enabled=True, transports=["http"]) + mcp.http.configure(token="from-code") + assert mcp.enabled is False + assert mcp.enabled_locked is True + assert mcp.transports == () # --no-mcp: nothing served + assert mcp.http.token == "from-cli" + + # the application's values are only defaults before start, however set + mcp.enabled = True + mcp.http.token = "from-attribute" + assert mcp.enabled is False + assert mcp.http.token == "from-cli" + + # once started every assignment applies, except what --no-mcp withholds + mcp._start() + mcp.http.token = "after-start" + assert mcp.http.token == "after-start" + mcp.enabled = True + assert mcp.enabled is False + + +def test_the_transports_from_the_command_line(monkeypatch): + monkeypatch.setenv("TRAME_ARGS", "--mcp-transport none") + mcp = new_server().mcp + assert mcp.transports == () + assert mcp.enabled is False # choosing transports doesn't turn access on + mcp.transports = ["http"] # the launch value wins until start + assert mcp.transports == () + + # only a starting value: once running, the app may open a transport + mcp._start() + mcp.transports = ["http"] + assert mcp.transports == ("http",) + + monkeypatch.setenv("TRAME_ARGS", "--mcp --mcp-transport http") + mcp = new_server().mcp + mcp.transports = [] + assert (mcp.enabled, mcp.transports) == (True, ("http",)) + + +@pytest.mark.parametrize("flags", ["--no-mcp-transport", "--no-mcp"]) +def test_no_transport_for_the_whole_run(monkeypatch, flags): + monkeypatch.setenv("TRAME_ARGS", flags) + mcp = new_server().mcp + assert mcp.transports == () + assert mcp.transports_locked is True + mcp.transports = ["http"] + mcp._start() + mcp.transports = ["http"] # held for the run + assert mcp.transports == () + mcp.transports = [] # withholding more is fine + if flags == "--no-mcp-transport": + mcp.enabled = True # agents in the process are still allowed + assert mcp.enabled is True + + +def test_only_withholding_launch_options_hold_after_start(monkeypatch): + monkeypatch.setenv("TRAME_ARGS", "--mcp") + mcp = new_server().mcp + assert mcp.enabled is True + assert mcp.enabled_locked is False + mcp.enabled = False # before start, the launch value wins + assert mcp.enabled is True + mcp._start() + mcp.enabled = False # turning on only set the starting value + assert mcp.enabled is False + + monkeypatch.setenv("TRAME_ARGS", "--no-mcp") + mcp = new_server().mcp + assert mcp.enabled_locked is True + mcp._start() + mcp.enabled = True # turning off holds for the run + assert mcp.enabled is False + + +def test_the_token_from_the_environment(monkeypatch): + monkeypatch.setenv("TRAME_MCP_HTTP_TOKEN", "from-env") + mcp = new_server().mcp + assert mcp.http.token == "from-env" + mcp.http.token = "from-code" + assert mcp.http.token == "from-env" + + # the command line comes first + monkeypatch.setenv("TRAME_ARGS", "--mcp-http-token from-cli") + assert new_server().mcp.http.token == "from-cli" + + +def test_launch_options_are_read_when_it_is_created(monkeypatch): + """Server.cli is the server's; the MCP reads it when first accessed.""" + server = new_server() + assert server.cli is server.cli + monkeypatch.setenv("TRAME_ARGS", "--mcp") + assert server.mcp.enabled is True + + +# ----------------------------------------------------------------------------- +# A running server +# ----------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_mcp_at_launch_turns_access_on(monkeypatch, capsys, caplog): + monkeypatch.setenv("TRAME_ARGS", "--mcp") + server = new_server() + server.mcp.tool(lambda: "pong", name="ping") + with caplog.at_level(logging.WARNING, logger="trame_server.mcp"): + await start(server) + try: + assert server.mcp.enabled is True + assert capsys.readouterr().out.count("MCP server available") == 1 + assert "exposes nothing" not in caplog.text + async with ( + aiohttp.ClientSession() as session, + session.post( + server.mcp.http.url, + json={"jsonrpc": "2.0", "id": 1, "method": "ping"}, + ) as response, + ): + assert response.status == 200 + finally: + await server.stop() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("create", [False, True]) +async def test_mcp_at_launch_with_nothing_exposed_warns(monkeypatch, caplog, create): + monkeypatch.setenv("TRAME_ARGS", "--mcp") + server = new_server() + if create: + server.mcp # noqa: B018 + with caplog.at_level(logging.WARNING, logger="trame_server.mcp"): + await start(server) + try: + assert "--mcp is given, but this app exposes nothing to agents" in caplog.text + finally: + await server.stop() + + +@pytest.mark.asyncio +async def test_the_url_names_the_bound_host(): + server = new_server() + mcp = server.mcp + await start(server, host="127.0.0.1") + try: + assert mcp.http.url == f"http://127.0.0.1:{server.port}/mcp" + assert mcp.http.token is None # loopback: no token needed + finally: + await server.stop() diff --git a/tests/test_mcp_endpoint.py b/tests/test_mcp_endpoint.py new file mode 100644 index 0000000..3d426e1 --- /dev/null +++ b/tests/test_mcp_endpoint.py @@ -0,0 +1,1169 @@ +"""The HTTP endpoint (``server.mcp.http``): the web-server-agnostic handler +with fake requests, then end to end on a real trame server.""" + +import asyncio +import base64 +import itertools +import json +import logging +import shlex +from types import SimpleNamespace + +import aiohttp +import pytest +import pytest_asyncio +from trame.app import get_server +from trame_common.decorators import mcp as mcp_decorators +from trame_common.obj.component import TrameComponent + +from trame_server.mcp import endpoint +from trame_server.mcp.endpoint import ( + HTTPRequest, + MCPEndpoint, + aiohttp_routes, + claude_command, + codex_command, + is_loopback, + mount, +) +from trame_server.mcp.protocol import ( + HEADER_MISMATCH, + INVALID_PARAMS, + INVALID_REQUEST, + METHOD_NOT_FOUND, + REQUEST_CANCELLED, + UNSUPPORTED_VERSION, +) + +LAUNCH_VARIABLES = ( + "TRAME_ARGS", + "TRAME_MCP_HTTP_TOKEN", + "TRAME_MCP_LOG", + "TRAME_MCP_DENY", +) +DISABLED = "MCP access is currently disabled by the application" + +_counter = itertools.count() + + +def rpc(method, params=None, rpc_id=1): + msg = {"jsonrpc": "2.0", "id": rpc_id, "method": method} + if params is not None: + msg["params"] = params + return msg + + +def modern(method, params=None, rpc_id=1, version="2026-07-28"): + """A request as clients of 2026-07-28 send it: its version and the + client's capabilities in _meta.""" + meta = { + "io.modelcontextprotocol/protocolVersion": version, + "io.modelcontextprotocol/clientCapabilities": {}, + } + return rpc(method, {**(params or {}), "_meta": meta}, rpc_id=rpc_id) + + +def mirrored(message): + """The headers a 2026-07-28 client sends with ``message``.""" + params = message["params"] + headers = { + "MCP-Protocol-Version": params["_meta"][ + "io.modelcontextprotocol/protocolVersion" + ], + "Mcp-Method": message["method"], + } + field = {"tools/call": "name", "prompts/get": "name", "resources/read": "uri"} + if message["method"] in field: + headers["Mcp-Name"] = params[field[message["method"]]] + return headers + + +def new_server(): + return get_server(f"test_mcp_endpoint_{next(_counter)}") + + +@pytest.fixture(autouse=True) +def _no_launch_options(monkeypatch): + for name in LAUNCH_VARIABLES: + monkeypatch.delenv(name, raising=False) + + +# ----------------------------------------------------------------------------- +# The handler, with fake requests +# ----------------------------------------------------------------------------- + + +def request( + message=None, + *, + headers=None, + host="localhost", + authority="localhost:8080", + content_type="application/json", + disconnected=lambda: False, +): + """A request as a web server adapter hands it to the endpoint.""" + if isinstance(message, bytes): + body = message + else: + body = json.dumps(rpc("ping") if message is None else message).encode() + all_headers = {"content-type": content_type} if content_type else {} + all_headers.update({k.lower(): v for k, v in (headers or {}).items()}) + + async def read(): + return body + + return HTTPRequest( + method="POST", + headers=all_headers, + host=host, + authority=authority, + read=read, + disconnected=disconnected, + ) + + +@pytest.fixture +def mcp(): + mcp = new_server().mcp + mcp.enabled = True + return mcp + + +@pytest.fixture +def handler(mcp): + return MCPEndpoint(mcp.server) + + +def test_is_loopback(): + for host in ( + "localhost", + "LOCALHOST", + "app.localhost", + "127.0.0.1", + "127.8.0.1", + "::1", + "[::1]", + ): + assert is_loopback(host), host + for host in ( + None, + "", + "0.0.0.0", + "::", + "192.168.1.5", + "example.com", + "localhost.evil.com", + ): + assert not is_loopback(host), host + + +@pytest.mark.asyncio +async def test_before_server_mcp_exists_it_answers_503(): + server = new_server() + child = server.create_child_server(prefix="child_") + for handler in (MCPEndpoint(server), MCPEndpoint(child)): + response = await handler.post(request()) + assert response.status == 503 + assert response.body == { + "jsonrpc": "2.0", + "id": None, + "error": {"code": -32000, "message": DISABLED}, + } + assert server._mcp is None # a request does not create it + + +@pytest.mark.asyncio +async def test_a_request_reaches_the_protocol(handler): + response = await handler.post(request(rpc("ping", rpc_id=4))) + assert response.status == 200 + assert response.body == {"jsonrpc": "2.0", "id": 4, "result": {}} + assert response.headers == {} + + # an invalid JSON-RPC message is answered in the body + response = await handler.post(request({"jsonrpc": "1.0", "id": 1})) + assert response.status == 200 + assert response.body["error"]["code"] == INVALID_REQUEST + + +@pytest.mark.asyncio +async def test_initialize_gives_each_client_a_session(handler): + first = await handler.post(request(rpc("initialize", {}))) + second = await handler.post(request(rpc("initialize", {}))) + assert first.status == second.status == 200 + session = first.headers["Mcp-Session-Id"] + assert len(session) >= 24 + assert second.headers["Mcp-Session-Id"] != session + + +@pytest.mark.asyncio +async def test_notifications_are_acknowledged_without_a_body(handler): + notification = {"jsonrpc": "2.0", "method": "notifications/initialized"} + response = await handler.post(request(notification)) + assert (response.status, response.body) == (202, None) + response = await handler.post(request({"jsonrpc": "2.0", "id": 9, "result": {}})) + assert response.status == 202 + + +@pytest.mark.asyncio +async def test_bodies_it_does_not_take(handler): + for body in (b"{not json", b"\xff\xfe"): + response = await handler.post(request(body)) + assert response.status == 400 + assert response.body["error"] == {"code": -32700, "message": "Parse error"} + response = await handler.post(request([rpc("ping")])) + assert response.status == 400 + assert response.body["error"] == { + "code": -32600, + "message": "JSON-RPC batching is not supported", + } + + +def test_get_and_delete_are_not_supported(): + response = MCPEndpoint.unsupported() + assert (response.status, response.body) == (405, None) + assert response.headers == {"Allow": "POST"} + + +@pytest.mark.asyncio +async def test_the_checks_run_in_order(handler, mcp): + evil = {"Origin": "https://evil.example"} + + # 1. Origin, before anything else + mcp.enabled = False + mcp.http.token = "s3cret" + response = await handler.post(request(headers=evil, content_type="text/plain")) + assert response.status == 403 + assert response.body["error"]["message"] == "Origin not allowed" + + # 2. the token + response = await handler.post(request(content_type="text/plain")) + assert response.status == 401 + assert response.headers == {"WWW-Authenticate": "Bearer"} + + # ... or without one, a loopback host + mcp.http.token = None + response = await handler.post( + request(host="myhost.lan", authority="myhost.lan:8080", content_type=None) + ) + assert response.status == 403 + assert "set a token" in response.body["error"]["message"] + + # 3. enabled + response = await handler.post(request(content_type="text/plain")) + assert response.status == 503 + + # 4. the content type + mcp.enabled = True + for content_type in ("text/plain", None, "application/jsonx"): + response = await handler.post(request(content_type=content_type)) + assert response.status == 415, content_type + response = await handler.post( + request(content_type="Application/JSON; charset=utf-8") + ) + assert response.status == 200 + + +@pytest.mark.asyncio +async def test_origins(handler, mcp): + async def status(origin, **kwargs): + response = await handler.post(request(headers={"Origin": origin}, **kwargs)) + return response.status + + assert await status("https://evil.example") == 403 + assert await status("null") == 403 # sandboxed frames, file:// + assert await status("not a url") == 403 + # Local browser tools (e.g. MCP Inspector) + assert await status("http://localhost:6274") == 200 + assert await status("http://127.0.0.1:6274") == 200 + assert await status("http://[::1]:6274") == 200 + + # Listed origins + mcp.http.allowed_origins = "https://trusted.example/" + assert mcp.http.allowed_origins == ("https://trusted.example",) + assert await status("https://trusted.example") == 200 + assert await status("https://other.example") == 403 + + # The page served next to the endpoint (its own origin) + mcp.http.token = "s3cret" + + async def from_page(origin): + response = await handler.post( + request( + headers={"Origin": origin, "Authorization": "Bearer s3cret"}, + host="app.example", + authority="app.example:8443", + ) + ) + return response.status + + assert await from_page("https://app.example:8443") == 200 + assert await from_page("https://app.example:9999") == 403 + + +@pytest.mark.asyncio +async def test_the_token(handler, mcp): + async def status(authorization=None, **kwargs): + headers = {"Authorization": authorization} if authorization else {} + return (await handler.post(request(headers=headers, **kwargs))).status + + mcp.http.token = "s3cret" + assert await status() == 401 + assert await status("Bearer nope") == 401 + assert await status("Basic s3cret") == 401 + assert await status("Bearer s3cret") == 200 + assert await status("bearer s3cret") == 200 + # with a token, the host name is no longer restricted to loopback + assert await status("Bearer s3cret", host="myhost.lan") == 200 + + mcp.http.token = None + assert await status() == 200 + assert await status(host="myhost.lan") == 403 + + +@pytest.mark.asyncio +async def test_a_client_that_disconnects_cancels_its_call(handler, mcp, monkeypatch): + monkeypatch.setattr(endpoint, "DISCONNECT_POLL", 0.01) + started, cleaned_up = asyncio.Event(), asyncio.Event() + gone = False + + @mcp.tool + async def wait_forever() -> str: + started.set() + try: + await asyncio.Event().wait() + finally: + cleaned_up.set() + return "done" + + @mcp.tool + def quick() -> str: + return "quick" + + posting = asyncio.ensure_future( + handler.post( + request( + rpc("tools/call", {"name": "wait_forever"}), + disconnected=lambda: gone, + ) + ) + ) + await asyncio.wait_for(started.wait(), 1) + gone = True + response = await asyncio.wait_for(posting, 1) + assert (response.status, response.body) == (499, None) + assert response.reason == "Client Closed Request" + assert cleaned_up.is_set() + + # The tool lock was released: the next call runs + response = await asyncio.wait_for( + handler.post(request(rpc("tools/call", {"name": "quick"}))), 1 + ) + assert response.body["result"]["content"][0]["text"] == "quick" + + +@pytest.mark.asyncio +async def test_a_session_scopes_cancellation(handler, mcp): + started = asyncio.Event() + + @mcp.tool + async def long_job() -> str: + started.set() + await asyncio.Event().wait() + return "done" + + session = {"Mcp-Session-Id": "abc"} + posting = asyncio.ensure_future( + handler.post( + request(rpc("tools/call", {"name": "long_job"}, rpc_id=3), headers=session) + ) + ) + await asyncio.wait_for(started.wait(), 1) + cancel = { + "jsonrpc": "2.0", + "method": "notifications/cancelled", + "params": {"requestId": 3}, + } + response = await handler.post(request(cancel, headers={"Mcp-Session-Id": "xyz"})) + assert response.status == 202 + await asyncio.sleep(0.01) + assert not posting.done() + await handler.post(request(cancel, headers=session)) + response = await asyncio.wait_for(posting, 1) + assert response.status == 200 + assert response.body["error"]["code"] == REQUEST_CANCELLED + + +def test_aiohttp_routes(): + routes = aiohttp_routes(MCPEndpoint(new_server()), "/api/mcp") + assert [(r.method, r.path) for r in routes] == [ + ("POST", "/api/mcp"), + ("GET", "/api/mcp"), + ("DELETE", "/api/mcp"), + ] + + +def test_mount_needs_an_aiohttp_application(caplog): + server = new_server() + advice = "serve trame_server.mcp.endpoint.MCPEndpoint yourself" + with caplog.at_level(logging.INFO, logger="trame_server.mcp"): + assert mount(server, SimpleNamespace(app=object())) is False + # only worth mentioning to an app that uses agents + assert advice not in caplog.text + server.mcp # noqa: B018 + assert mount(server, SimpleNamespace()) is False + assert advice in caplog.text + + +def test_the_claude_command_adds_the_server_for_one_session(): + assert claude_command("app", "http://localhost:8080/mcp") == ( + "claude --mcp-config " + """'{"mcpServers":{"app":{"type":"http","url":"http://localhost:8080/mcp"}}}'""" + ) + assert claude_command("app", "http://h:8080/mcp", "s3cret").endswith( + """"headers":{"Authorization":"Bearer s3cret"}}}}'""" + ) + # The JSON stays one shell word whatever the name holds + assert shlex.split(claude_command("it's", "http://h/mcp"))[-1] == ( + """{"mcpServers":{"it's":{"type":"http","url":"http://h/mcp"}}}""" + ) + + +def test_the_codex_command_adds_the_server_for_one_session(): + assert codex_command("app", "http://localhost:8080/mcp") == ( + """codex -c 'mcp_servers.app.url="http://localhost:8080/mcp"'""" + ) + assert codex_command("app", "http://h:8080/mcp", "s3cret") == ( + """codex -c 'mcp_servers.app.url="http://h:8080/mcp"' """ + """-c 'mcp_servers.app.http_headers={Authorization="Bearer s3cret"}'""" + ) + + +# ----------------------------------------------------------------------------- +# End to end: a real trame server, MCP over HTTP +# ----------------------------------------------------------------------------- + + +async def start(server, **kwargs): + server.start( + exec_mode="task", + port=0, + open_browser=False, + show_connection_info=False, + disable_logging=True, + **kwargs, + ) + await server.ready + + +async def post(session, server, message=None, path="/mcp", **kwargs): + url = f"http://localhost:{server.port}{path}" + if "data" not in kwargs: + kwargs.setdefault("json", message or rpc("ping")) + async with session.post(url, **kwargs) as resp: + body = await resp.json(content_type=None) if resp.status != 202 else None + return resp.status, body + + +async def status_of(session, server, path="/mcp", method="POST"): + url = f"http://localhost:{server.port}{path}" + kwargs = {"json": rpc("ping")} if method == "POST" else {} + async with session.request(method, url, **kwargs) as resp: + return resp.status, resp.headers.get("Allow", "") + + +@pytest_asyncio.fixture +async def running_server(): + server = new_server() + mcp = server.mcp + mcp.enabled = True + + server.state.resolution = 6 + server.state.mcp_expose("resolution", type=int) + + @mcp.tool + def add(a: int, b: int) -> int: + """Add two numbers.""" + return a + b + + await start(server) + yield server, mcp + await server.stop() + + +@pytest.mark.asyncio +async def test_http_roundtrip(running_server): + server, mcp = running_server + url = f"http://localhost:{server.port}/mcp" + assert mcp.http.url == url + + async with aiohttp.ClientSession() as session: + # initialize: the preferred version, or the client's when it is known + status, body = await post(session, server, rpc("initialize", {})) + assert status == 200 + assert body["result"]["protocolVersion"] == "2025-11-25" + _, body = await post( + session, server, rpc("initialize", {"protocolVersion": "2025-06-18"}) + ) + assert body["result"]["protocolVersion"] == "2025-06-18" + + # initialized notification -> 202, empty + status, _ = await post( + session, server, {"jsonrpc": "2.0", "method": "notifications/initialized"} + ) + assert status == 202 + + # tools/list + _, body = await post(session, server, rpc("tools/list")) + assert {"add", "object_get", "object_update"} <= { + t["name"] for t in body["result"]["tools"] + } + + # tools/call -> touches real trame state + update = {"target": "state", "values": {"resolution": 24}} + _, body = await post( + session, + server, + rpc("tools/call", {"name": "object_update", "arguments": update}), + ) + assert body["result"].get("isError") is None + assert server.state.resolution == 24 + + # schema-checked: a string never reaches the state + update = {"target": "state", "values": {"resolution": "x"}} + _, body = await post( + session, + server, + rpc("tools/call", {"name": "object_update", "arguments": update}), + ) + assert body["result"]["isError"] is True + assert server.state.resolution == 24 + + # GET and DELETE are not supported (stateless server) + assert await status_of(session, server, method="GET") == (405, "POST") + assert await status_of(session, server, method="DELETE") == (405, "POST") + + # batch rejected + status, _ = await post(session, server, [rpc("ping")]) + assert status == 400 + + # malformed json + status, _ = await post( + session, + server, + data=b"{not json", + headers={"Content-Type": "application/json"}, + ) + assert status == 400 + + +@pytest.mark.asyncio +async def test_cross_site_requests_are_rejected(running_server): + server, mcp = running_server + attack = rpc( + "tools/call", + { + "name": "object_update", + "arguments": {"target": "state", "values": {"resolution": 99}}, + }, + ) + + async with aiohttp.ClientSession() as session: + # A page on another site: browsers attach its Origin + status, _ = await post( + session, server, attack, headers={"Origin": "https://evil.example"} + ) + assert status == 403 + status, _ = await post(session, server, attack, headers={"Origin": "null"}) + assert status == 403 + + # "Simple" request a browser sends without CORS preflight + status, _ = await post( + session, server, data=b"{}", headers={"Content-Type": "text/plain"} + ) + assert status == 415 + + # DNS rebinding: attacker domain resolved to 127.0.0.1, same-origin to + # the browser + rebound = f"evil.example:{server.port}" + status, body = await post( + session, + server, + attack, + headers={"Host": rebound, "Origin": f"http://{rebound}"}, + ) + assert status == 403 + assert "set a token" in body["error"]["message"] + + assert server.state.resolution == 6 + + # Local browser tools (e.g. MCP Inspector) and listed origins are fine + status, _ = await post( + session, server, headers={"Origin": "http://localhost:6274"} + ) + assert status == 200 + mcp.http.allowed_origins = ["https://trusted.example/"] + status, _ = await post( + session, server, headers={"Origin": "https://trusted.example"} + ) + assert status == 200 + + +@pytest.mark.asyncio +async def test_a_2026_client(running_server): + """No handshake and no session: each request names its version.""" + server, _ = running_server + url = f"http://localhost:{server.port}/mcp" + call = modern("tools/call", {"name": "add", "arguments": {"a": 2, "b": 3}}) + async with aiohttp.ClientSession() as session: + async with session.post(url, json=call, headers=mirrored(call)) as resp: + assert resp.status == 200 + assert "Mcp-Session-Id" not in resp.headers + body = await resp.json() + assert body["result"]["content"] == [{"type": "text", "text": "5"}] + assert body["result"]["resultType"] == "complete" + + discover = modern("server/discover") + status, body = await post(session, server, discover, headers=mirrored(discover)) + assert status == 200 + assert body["result"]["supportedVersions"][0] == "2026-07-28" + + +@pytest.mark.asyncio +async def test_a_2026_request_mirrors_its_body_in_headers(running_server): + server, _ = running_server + call = modern("tools/call", {"name": "add", "arguments": {"a": 1}}) + good = mirrored(call) + + async def send(headers, message=call): + async with aiohttp.ClientSession() as session: + return await post(session, server, message, headers=headers) + + for missing in good: + status, body = await send({k: v for k, v in good.items() if k != missing}) + assert status == 400 + assert body["id"] == 1 + assert body["error"]["code"] == HEADER_MISMATCH + assert f"the {missing} header is missing" in body["error"]["message"] + + for header in ("Mcp-Name", "Mcp-Method", "MCP-Protocol-Version"): + status, body = await send({**good, header: "other"}) + assert status == 400 + assert "does not match body value" in body["error"]["message"] + + # A name outside plain ASCII is sent in the Base64 form + encoded = "=?base64?" + base64.b64encode(b"add").decode() + "?=" + assert (await send({**good, "Mcp-Name": encoded}))[0] == 200 + assert (await send({**good, "Mcp-Name": "=?base64?not base64?="}))[0] == 400 + + # Headers of 2026-07-28 on a body that doesn't name its version + status, body = await send( + {"MCP-Protocol-Version": "2026-07-28", "Mcp-Method": "tools/list"}, + message=rpc("tools/list"), + ) + assert (status, body["error"]["code"]) == (400, HEADER_MISMATCH) + + +@pytest.mark.asyncio +async def test_a_2026_request_the_server_cannot_serve(running_server): + server, _ = running_server + + async def send(message): + async with aiohttp.ClientSession() as session: + return await post(session, server, message, headers=mirrored(message)) + + status, body = await send(modern("tools/list", version="2099-01-01")) + assert (status, body["error"]["code"]) == (400, UNSUPPORTED_VERSION) + assert body["error"]["data"]["requested"] == "2099-01-01" + + message = modern("tools/list") + del message["params"]["_meta"]["io.modelcontextprotocol/clientCapabilities"] + status, body = await send(message) + assert (status, body["error"]["code"]) == (400, INVALID_PARAMS) + + for method in ("ping", "initialize", "subscriptions/listen"): + status, body = await send(modern(method)) + assert (status, body["id"], body["error"]["code"]) == (404, 1, METHOD_NOT_FOUND) + + +@pytest.mark.asyncio +async def test_token_auth(running_server): + server, mcp = running_server + mcp.http.token = "s3cret" + + async with aiohttp.ClientSession() as session: + url = f"http://localhost:{server.port}/mcp" + async with session.post(url, json=rpc("ping")) as resp: + assert resp.status == 401 + assert resp.headers["WWW-Authenticate"] == "Bearer" + status, _ = await post( + session, server, headers={"Authorization": "Bearer nope"} + ) + assert status == 401 + status, _ = await post( + session, server, headers={"Authorization": "bearer s3cret"} + ) + assert status == 200 + + # with a token, the host name is no longer restricted to loopback + status, _ = await post( + session, + server, + headers={ + "Authorization": "Bearer s3cret", + "Host": f"myhost.lan:{server.port}", + }, + ) + assert status == 200 + + mcp.http.token = None + status, _ = await post(session, server) + assert status == 200 + + +@pytest.mark.asyncio +async def test_runtime_enable_switch(running_server): + server, mcp = running_server + + async with aiohttp.ClientSession() as session: + mcp.enabled = False + status, body = await post(session, server) + assert status == 503 + assert body == { + "jsonrpc": "2.0", + "id": None, + "error": {"code": -32000, "message": DISABLED}, + } + + mcp.enabled = True + status, _ = await post(session, server) + assert status == 200 + + +@pytest.mark.asyncio +async def test_path_is_fixed_once_started(running_server): + _, mcp = running_server + with pytest.raises(RuntimeError, match="started"): + mcp.http.path = "/other" + assert mcp.http.path == "/mcp" + + +@pytest.mark.asyncio +async def test_initialize_gives_a_session_that_scopes_cancellation(running_server): + server, mcp = running_server + url = f"http://localhost:{server.port}/mcp" + started, cleaned_up = asyncio.Event(), asyncio.Event() + + @mcp.tool + async def long_job() -> str: + started.set() + try: + await asyncio.Event().wait() # until cancelled + finally: + cleaned_up.set() + return "done" + + async with aiohttp.ClientSession() as session: + async with session.post(url, json=rpc("initialize", {})) as resp: + session_id = resp.headers["Mcp-Session-Id"] + async with session.post(url, json=rpc("initialize", {})) as resp: + assert resp.headers["Mcp-Session-Id"] != session_id # one per client + # other responses carry none + async with session.post(url, json=rpc("ping")) as resp: + assert "Mcp-Session-Id" not in resp.headers + + headers = {"Mcp-Session-Id": session_id} + job = asyncio.ensure_future( + post( + session, + server, + rpc("tools/call", {"name": "long_job"}, rpc_id=3), + headers=headers, + ) + ) + await asyncio.wait_for(started.wait(), 2) + + cancel = { + "jsonrpc": "2.0", + "method": "notifications/cancelled", + "params": {"requestId": 3, "reason": "timed out"}, + } + # Another session's cancel leaves it running + status, _ = await post( + session, server, cancel, headers={"Mcp-Session-Id": "other"} + ) + assert status == 202 + await asyncio.sleep(0.05) + assert not job.done() + + status, _ = await post(session, server, cancel, headers=headers) + assert status == 202 + status, body = await asyncio.wait_for(job, 2) + assert status == 200 + assert body["id"] == 3 + assert body["error"] == { + "code": REQUEST_CANCELLED, + "message": "Request cancelled: timed out", + } + assert cleaned_up.is_set() + + +@pytest.mark.asyncio +async def test_a_client_that_disconnects_over_http_cancels_its_call( + running_server, monkeypatch +): + monkeypatch.setattr(endpoint, "DISCONNECT_POLL", 0.02) + server, mcp = running_server + url = f"http://localhost:{server.port}/mcp" + started, cleaned_up = asyncio.Event(), asyncio.Event() + + @mcp.tool + async def wait_forever() -> str: + started.set() + try: + await asyncio.Event().wait() + finally: + cleaned_up.set() + return "done" + + async with aiohttp.ClientSession() as session: + # A client giving up: its request times out and the connection closes + with pytest.raises(asyncio.TimeoutError): + await session.post( + url, + json=rpc("tools/call", {"name": "wait_forever"}), + timeout=aiohttp.ClientTimeout(total=0.3), + ) + assert started.is_set() + await asyncio.wait_for(cleaned_up.wait(), 2) + + # The tool lock was released: the next call runs + status, body = await post( + session, + server, + rpc("tools/call", {"name": "add", "arguments": {"a": 1, "b": 2}}), + ) + assert status == 200 + assert body["result"]["content"][0]["text"] == "3" + + +# ----------------------------------------------------------------------------- +# The route at start, and the banner +# ----------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_the_route_exists_before_server_mcp_does(capsys): + """Every app mounts the route at start; it answers 503 until the app + creates server.mcp and turns it on.""" + server = new_server() + await start(server) + try: + assert server._mcp is None + assert server._mcp_endpoint_mounted is True + async with aiohttp.ClientSession() as session: + status, body = await post(session, server) + assert status == 503 + assert body["error"]["message"] == DISABLED + assert server._mcp is None + + mcp = server.mcp # created once the server runs + assert mcp.http.url == f"http://localhost:{server.port}/mcp" + with pytest.raises(RuntimeError, match="started"): + mcp.http.path = "/other" + status, _ = await post(session, server) + assert status == 503 + assert "MCP server available" not in capsys.readouterr().out + + mcp.enabled = True + assert f"MCP server available at {mcp.http.url}" in capsys.readouterr().out + status, body = await post(session, server) + assert (status, body["result"]) == (200, {}) + finally: + await server.stop() + + +@pytest.mark.asyncio +async def test_off_until_enabled_then_announced(capsys, caplog): + """Declared but off: 503 and no banner; once running, the switch applies + and the banner prints when access goes live, once.""" + server = new_server() + mcp = server.mcp + + @mcp.tool + def hello() -> str: + """Say hi.""" + return "hi" + + with caplog.at_level(logging.INFO, logger="trame_server.mcp"): + await start(server) + try: + assert f"{mcp.http.url} is mounted, but agent access is off" in caplog.text + assert mcp.enabled is False + assert "MCP server available" not in capsys.readouterr().out + async with aiohttp.ClientSession() as session: + status, _ = await post(session, server) + assert status == 503 + + mcp.enabled = True + out = capsys.readouterr().out + assert f"MCP server available at {mcp.http.url}" in out + assert claude_command(mcp.name, mcp.http.url) in out + assert codex_command(mcp.name, mcp.http.url) in out + status, _ = await post(session, server) + assert status == 200 + + mcp.enabled = False + mcp.enabled = True + assert "MCP server available" not in capsys.readouterr().out # once + finally: + await server.stop() + + +@pytest.mark.asyncio +async def test_enabled_before_start_is_announced_when_ready(capsys): + server = new_server() + mcp = server.mcp + mcp.configure(name="viewer", enabled=True) + mcp.http.token = "s3cret" + assert "MCP server available" not in capsys.readouterr().out + + await start(server) + try: + out = capsys.readouterr().out + assert out.count("MCP server available") == 1 + # a token the user set is shown as a placeholder + assert claude_command("viewer", mcp.http.url, "") in out + assert "s3cret" not in out + finally: + await server.stop() + + +@pytest.mark.asyncio +async def test_the_banner_follows_the_connection_information(capsys): + """Apps started in main mode print the connection information when ready; + the banner follows it, once.""" + from trame_server.utils.server import print_informations # noqa: PLC0415 + + server = new_server() + server.mcp.enabled = True + server.controller.on_server_ready.add(lambda **_: print_informations(server)) + await start(server) + try: + out = capsys.readouterr().out + assert out.count("MCP server available") == 1 + assert out.index("App running at:") < out.index("MCP server available") + finally: + await server.stop() + + +@pytest.mark.asyncio +async def test_no_mcp_serves_no_endpoint(monkeypatch, capsys): + monkeypatch.setenv("TRAME_ARGS", "--no-mcp") + server = new_server() + mcp = server.mcp + mcp.enabled = True + assert mcp.enabled_locked is True + assert mcp.enabled is False + + await start(server) + try: + assert server._mcp_endpoint_mounted is False + mcp.enabled = True # a UI switch, once running: still off + assert mcp.enabled is False + assert "MCP server available" not in capsys.readouterr().out + async with aiohttp.ClientSession() as session: + # No MCP route: only trame's static files answer + status, allow = await status_of(session, server, method="GET") + assert status == 404 + status, allow = await status_of(session, server) + assert status in (404, 405) + assert "POST" not in allow + finally: + await server.stop() + + +@pytest.mark.asyncio +async def test_no_mcp_without_server_mcp(monkeypatch): + monkeypatch.setenv("TRAME_ARGS", "--no-mcp") + server = new_server() + await start(server) + try: + assert server._mcp_endpoint_mounted is False + async with aiohttp.ClientSession() as session: + assert (await status_of(session, server, method="GET"))[0] == 404 + assert server.mcp.enabled_locked is True + finally: + await server.stop() + + +@pytest.mark.asyncio +async def test_the_http_transport_switched_at_runtime(capsys, caplog): + """transports=[]: the route is there but answers 503 and no banner + prints, while agents in the process reach what the app declares. Adding + http later serves it, and the banner prints then.""" + server = new_server() + mcp = server.mcp + mcp.configure(enabled=True, transports=[]) + + @mcp.tool + def hello() -> str: + """Say hi.""" + return "hi" + + with caplog.at_level(logging.INFO, logger="trame_server.mcp"): + await start(server) + try: + assert server._mcp_endpoint_mounted is True + assert "MCP server available" not in capsys.readouterr().out + assert ( + f"{mcp.http.url} is mounted, but the HTTP transport is off" in caplog.text + ) + response = await mcp.protocol.handle(rpc("tools/call", {"name": "hello"})) + assert response["result"]["content"] == [{"type": "text", "text": "hi"}] + async with aiohttp.ClientSession() as session: + status, body = await post(session, server) + assert status == 503 + assert body["error"]["message"] == ( + "The application turned its HTTP transport off" + ) + + mcp.transports = ["http"] + assert f"MCP server available at {mcp.http.url}" in capsys.readouterr().out + assert (await post(session, server))[0] == 200 + + mcp.transports = [] + assert (await post(session, server))[0] == 503 + mcp.transports = ["http"] + assert "MCP server available" not in capsys.readouterr().out # once + finally: + await server.stop() + + +@pytest.mark.asyncio +async def test_no_transport_at_start_keeps_the_route(monkeypatch): + """--mcp-transport none only sets the starting value.""" + monkeypatch.setenv("TRAME_ARGS", "--mcp-transport none") + server = new_server() + await start(server) + try: + assert server._mcp_endpoint_mounted is True + mcp = server.mcp + assert (mcp.transports, mcp.transports_locked) == ((), False) + mcp.configure(enabled=True, transports=["http"]) + async with aiohttp.ClientSession() as session: + assert (await post(session, server))[0] == 200 + finally: + await server.stop() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("create", [False, True]) +async def test_no_mcp_transport_serves_no_endpoint(monkeypatch, create): + monkeypatch.setenv("TRAME_ARGS", "--no-mcp-transport") + server = new_server() + if create: + server.mcp # noqa: B018 + await start(server) + try: + assert server._mcp_endpoint_mounted is False + mcp = server.mcp + assert mcp.transports_locked is True + async with aiohttp.ClientSession() as session: + assert (await status_of(session, server, method="GET"))[0] == 404 + mcp.enabled = True # agents in the process only + response = await mcp.protocol.handle(rpc("tools/list")) + assert "result" in response + finally: + await server.stop() + + +@pytest.mark.asyncio +async def test_a_path_set_before_start(capsys): + server = new_server() + mcp = server.mcp + mcp.configure(enabled=True) + mcp.http.configure(path="api/mcp") + assert mcp.http.path == "/api/mcp" + + await start(server) + try: + assert mcp.http.url == f"http://localhost:{server.port}/api/mcp" + assert mcp.http.url in capsys.readouterr().out + async with aiohttp.ClientSession() as session: + status, _ = await post(session, server, path="/api/mcp") + assert status == 200 + status, allow = await status_of(session, server, path="/mcp") + assert "POST" not in allow + assert status != 200 + finally: + await server.stop() + + +@pytest.mark.asyncio +async def test_a_component_does_not_turn_mcp_on(capsys): + """A component declaring methods for agents serves nothing until the app + turns agent access on; building it does not even create server.mcp.""" + + class EndpointViewer(TrameComponent): + @mcp_decorators.method + def reset_camera(self) -> str: + """Reset the camera.""" + return "reset" + + server = new_server() + viewer = EndpointViewer(server) + assert server._mcp is None + await start(server) + try: + assert "MCP server available" not in capsys.readouterr().out + async with aiohttp.ClientSession() as session: + status, _ = await post(session, server, rpc("tools/list")) + assert status == 503 + finally: + await server.stop() + del viewer + + +@pytest.mark.asyncio +async def test_token_generated_when_listening_beyond_localhost(capsys): + server = new_server() + mcp = server.mcp + mcp.enabled = True + assert mcp.http.token is None + + await start(server, host="0.0.0.0") + try: + assert mcp.http.token + assert mcp.http.url == f"http://localhost:{server.port}/mcp" + out = capsys.readouterr().out + assert f"Bearer {mcp.http.token}" in out + assert "token generated because the server listens beyond localhost" in out + + async with aiohttp.ClientSession() as session: + status, _ = await post(session, server) + assert status == 401 + status, _ = await post( + session, server, headers={"Authorization": f"Bearer {mcp.http.token}"} + ) + assert status == 200 + finally: + await server.stop() + + +@pytest.mark.asyncio +async def test_a_token_set_after_one_was_generated_is_not_shown(capsys): + server = new_server() + mcp = server.mcp + await start(server, host="0.0.0.0") + try: + assert mcp.http.token # generated: the server listens beyond localhost + mcp.http.token = "mine" + mcp.enabled = True + out = capsys.readouterr().out + assert "MCP server available" in out + assert "Bearer " in out + assert "mine" not in out + assert "token generated" not in out + finally: + await server.stop() diff --git a/tests/test_mcp_objects.py b/tests/test_mcp_objects.py new file mode 100644 index 0000000..b311e41 --- /dev/null +++ b/tests/test_mcp_objects.py @@ -0,0 +1,1461 @@ +"""The object layer with components and the state object (models are tested +in trame-dataclass).""" + +import asyncio +import gc +import itertools +import json +import logging +from typing import Annotated, Literal + +import pytest +from trame.app import TrameComponent, get_server +from trame_common.decorators import mcp +from trame_common.obj.mcp import MCPExposable + +from trame_server.mcp.objects import OBJECT_TOOLS + +_counter = itertools.count() + + +@pytest.fixture +def server(): + server = get_server(f"test_mcp_objects_{next(_counter)}") + server.mcp.enabled = True + return server + + +async def call(server, name, **arguments): + """``(value, is_error)`` of a tools/call, as an agent reads it.""" + response = await server.mcp.protocol.handle( + { + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": {"name": name, "arguments": arguments}, + } + ) + result = response["result"] + content = result["content"] + text = content[0]["text"] if 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 + + +async def fails(server, name, **arguments): + value, error = await call(server, name, **arguments) + assert 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]} + + +# ----------------------------------------------------------------------------- +# Classes (names are process-wide for agents: keep them unique in the suite) +# ----------------------------------------------------------------------------- + + +class Viewport(TrameComponent): + """A 3D view of the scene. It shows what is loaded.""" + + def __init__(self, server, background="black"): + super().__init__(server) + self._background = background + self.calls = [] + + @mcp.method + def set_view_up(self, axis: Literal["x", "y", "z"]) -> str: + """Orient the camera so this axis points up.""" + self.calls.append(axis) + return f"up {axis}" + + @mcp.method(read_only=True) + def describe_view(self) -> dict: + """What the view shows.""" + return {"background": self._background} + + @mcp.method + async def wait(self, seconds: Annotated[float, {"minimum": 0}] = 0) -> str: + return f"waited {seconds}" + + @mcp.method(read_only=True) + def echo(self, value: float | int | str) -> str: # noqa: PYI041 (the point) + return f"{type(value).__name__} {value}" + + @mcp.property + def background(self) -> Annotated[str, "Background color"]: + return self._background + + @background.setter + def background(self, value): + if value == "boom": + msg = "cannot" + raise RuntimeError(msg) + self._background = value + + @mcp.property + def size(self) -> int: + """Pixels across.""" + return 300 + + +class Desk(TrameComponent): + """Holds the viewports.""" + + def __init__(self, server): + super().__init__(server) + self.views = [] + + @mcp.property + def active(self) -> Viewport | None: + return self.views[0] if self.views else None + + @mcp.method + def pick(self, view: Viewport) -> Viewport: + """Make a viewport the active one.""" + self.views.remove(view) + self.views.insert(0, view) + return view + + +class Plain(TrameComponent): + def do(self): ... + + +class Locked(Viewport, mcp=False): + pass + + +# ----------------------------------------------------------------------------- +# Exposure +# ----------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_reached_once_named(server): + view = Viewport(server) + assert "object_get" not in server.mcp._tools # built is not exposed + + view.mcp_expose("view") + assert set(OBJECT_TOOLS) <= set(server.mcp._tools) + + assert await ok(server, "object_list") == { + "objects": {"view": {"type": "Viewport"}} + } + index = await ok(server, "object_types") + assert index["types"]["Viewport"] == { + "about": "A 3D view of the scene.", + "properties": ["background", "size"], + "methods": ["set_view_up", "describe_view", "wait", "echo"], # as defined + } + + value = await ok(server, "object_get", target=view._id) + assert value == { + "result": {"id": view._id}, + "objects": { + view._id: { + "type": "Viewport", + "names": ["view"], + "values": {"background": "black", "size": 300}, + } + }, + } + + +@pytest.mark.asyncio +async def test_names(server): + left, right = Viewport(server), Viewport(server) + left.mcp_expose("left") + left.mcp_expose("main") + right.mcp_expose("right") + + value = read(await ok(server, "object_get", target="main")) + assert value["id"] == left._id + assert value["names"] == ["left", "main"] + assert (await ok(server, "object_list"))["objects"] == { + "left": {"type": "Viewport"}, + "main": {"type": "Viewport"}, + "right": {"type": "Viewport"}, + } + + with pytest.raises(ValueError, match="starts with a letter"): + left.mcp_expose("1st") + with pytest.raises(ValueError, match="reserved"): + left.mcp_expose("state") + with pytest.raises(TypeError, match="declares nothing"): + server.mcp.expose_object(Plain(server), "plain") + with pytest.raises(TypeError, match="takes a component"): + server.mcp.expose_object(3, "three") + + server.mcp.unexpose_object("main") + assert read(await ok(server, "object_get", target="left"))["names"] == ["left"] + with pytest.raises(KeyError): + server.mcp.unexpose_object("main") + + message = await fails(server, "object_get", target="viewer") + assert message == ( + 'Unknown object "viewer". Objects: left, right (or an object id); ' + "object_list() lists them all." + ) + + +@pytest.mark.asyncio +async def test_getter_follows_reassignment(server): + first, second = Viewport(server, "red"), Viewport(server, "blue") + holder = {"active": first} + server.mcp.expose_object(lambda: holder["active"], "active_viewer") + + value = await ok(server, "object_get", target="active_viewer") + assert value["result"] == {"id": first._id} + holder["active"] = second + value = await ok(server, "object_get", target="active_viewer") + assert value["result"] == {"id": second._id} + holder["active"] = None + message = await fails(server, "object_get", target="active_viewer") + assert "currently empty" in message + + +@pytest.mark.asyncio +async def test_unexpose_and_unbind(server): + view = Viewport(server) + view.mcp_expose("view") + view.mcp_unexpose() + assert "Unknown object" in await fails(server, "object_get", target=view._id) + assert "Unknown object" in await fails(server, "object_get", target="view") + + other = Viewport(server) + other.mcp_expose("other") + await ok(server, "object_get", target=other._id) + other._unbind_annotated_methods() + assert "Unknown object" in await fails(server, "object_get", target=other._id) + + +@pytest.mark.asyncio +async def test_components_are_held_weakly(server): + view = Viewport(server) + view.mcp_expose("view") + _id = view._id + del view + gc.collect() + assert await ok(server, "object_list") == {"objects": {"view": {"type": None}}} + assert "Unknown object" in await fails(server, "object_get", target=_id) + Viewport(server).mcp_expose("view") # the name is free again + + +class Lens(MCPExposable, mcp={"description": "A camera lens"}): + """A plain class: it mixes in MCPExposable, with no server of its own.""" + + def __init__(self): + self._zoom = 1.0 + + @mcp.property + def zoom(self) -> float: + return self._zoom + + @zoom.setter + def zoom(self, value: float): + self._zoom = value + + @mcp.method + def reset(self) -> str: + """Zoom back out.""" + self._zoom = 1.0 + return "Reset" + + +@pytest.mark.asyncio +async def test_any_class_can_mix_in_exposable(server): + server.mcp.expose_object(Lens(), "lens") + gc.collect() # held: unlike a component's, nothing else keeps it alive + await ok(server, "object_update", target="lens", values={"zoom": 2.5}) + assert read(await ok(server, "object_get", target="lens"))["values"] == { + "zoom": 2.5 + } + assert await ok(server, "object_call", target="lens", method="reset") == "Reset" + described = (await ok(server, "object_types", type_names=["Lens"]))["types"] + assert described["Lens"]["description"] == "A camera lens" + + class Leaky(MCPExposable, mcp={"methods": ["mcp_expose"]}): + pass + + with pytest.raises(ValueError, match="part of trame's API"): + server.mcp.describe(Leaky) + + +class Shelf(TrameComponent): + """Holds viewports in nested containers.""" + + def __init__(self, server): + super().__init__(server) + self._rows = {"top": [Viewport(server), Viewport(server)]} + + @mcp.property + def rows(self) -> dict[str, list[Viewport]]: + return self._rows + + @mcp.method(read_only=True) + def distinct(self) -> set[Viewport]: + """The viewports, once each.""" + return set(self._rows["top"]) + + +@pytest.mark.asyncio +async def test_objects_nested_in_containers_are_reached(server): + shelf = Shelf(server) + shelf.mcp_expose("shelf") + top = [view._id for view in shelf._rows["top"]] + reply = await ok(server, "object_get", target="shelf", depth=2) + rows = read(reply)["values"]["rows"] + assert sorted(ref["id"] for ref in rows["top"]) == sorted(top) + assert set(top) <= reply["objects"].keys() + assert read(await ok(server, "object_get", target=top[0]))["id"] == top[0] + + # A set of objects comes back as references too, not as text + reply = await ok(server, "object_call", target="shelf", method="distinct") + assert sorted(ref["id"] for ref in reply["result"]) == sorted(top) + + +@pytest.mark.asyncio +async def test_a_name_exposes_one_object(server, caplog): + view, other = Viewport(server), Viewport(server) + with pytest.raises(TypeError, match="each under its own name"): + server.mcp.expose_object([view, other], "views") + + view.mcp_expose("view", description="The main view.") + view.mcp_expose("view", read_only=True) # the same object: new options + listed = (await ok(server, "object_list"))["objects"] + assert listed["view"] == {"type": "Viewport", "read_only": True} + with pytest.raises(ValueError, match=f"'view' already exposes Viewport {view._id}"): + other.mcp_expose("view") + + def active(): + return other + + server.mcp.expose_object(active, "active") + server.mcp.expose_object(active, "active", read_only=True) + with pytest.raises(ValueError, match="'active' already exposes a function"): + server.mcp.expose_object(lambda: other, "active") + + # Withdrawn, the name is free + view.mcp_unexpose() + other.mcp_expose("view") + assert read(await ok(server, "object_get", target="view"))["id"] == other._id + + # A function returning something else exposes nothing + server.mcp.expose_object(lambda: [view], "listed") + listed = (await ok(server, "object_list"))["objects"] + assert listed["listed"] == {"type": None} + assert "'listed' exposes a list" in caplog.text + + +@pytest.mark.asyncio +async def test_descriptions(server): + left, right = Viewport(server), Viewport(server) + left.mcp_expose("left", description="The view on the left. It shows the input.") + right.mcp_expose("right", description="The view on the right.", read_only=True) + listed = (await ok(server, "object_list"))["objects"] + assert listed == { + "left": { + "type": "Viewport", + "description": "The view on the left. It shows the input.", + }, + "right": { + "type": "Viewport", + "read_only": True, + "description": "The view on the right.", + }, + } + + # The connection instructions keep the first sentence + response = await server.mcp.protocol.handle( + {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}} + ) + assert ( + "Objects: left (Viewport: The view on the left.), right (Viewport, " + "read-only: The view on the right.); object_list() lists them" + in response["result"]["instructions"] + ) + + with pytest.raises(TypeError, match="description must be a string"): + left.mcp_expose("main", description=3) + + +# ----------------------------------------------------------------------------- +# Read-only +# ----------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_read_only_component(server): + view = Viewport(server) + view.mcp_expose("view", read_only=True) + + for target in ("view", view._id): + value = read(await ok(server, "object_get", target=target)) + assert value["read_only"] is True + message = await fails( + server, + "object_call", + target=target, + method="set_view_up", + args={"axis": "y"}, + ) + assert "read-only" in message + assert "describe_view" in message + assert await ok(server, "object_call", target=target, method="describe_view") + message = await fails( + server, "object_update", target=target, values={"background": "white"} + ) + assert "read-only" in message + assert view.calls == [] + + # A writable name wins + view.mcp_expose("editable") + assert await ok( + server, "object_call", target=view._id, method="set_view_up", args={"axis": "x"} + ) + + +@pytest.mark.asyncio +async def test_read_only_function_covers_its_result(server): + views = [Viewport(server), Viewport(server)] + server.mcp.expose_object(lambda: views[0], "first", read_only=True) + + message = await fails( + server, + "object_call", + target=views[0]._id, + method="set_view_up", + args={"axis": "y"}, + ) + assert "read-only" in message + listed = (await ok(server, "object_list"))["objects"] + assert listed == {"first": {"type": "Viewport", "read_only": True}} + assert "Unknown object" in await fails( + server, "object_get", target=views[1]._id + ) # not what it returns + + # A writable name wins + views[0].mcp_expose("editable") + assert await ok( + server, + "object_call", + target=views[0]._id, + method="set_view_up", + args={"axis": "x"}, + ) + + +# ----------------------------------------------------------------------------- +# object_list and object_types +# ----------------------------------------------------------------------------- + + +class Base3D(TrameComponent): + """The base of 3D things.""" + + @mcp.method + def reset(self) -> None: + """Back to defaults.""" + + @mcp.property + def shape(self) -> str: + return "base" + + +class Derived3D(Base3D): + @mcp.method + def spin(self, turns: int = 1) -> None: + """Spin it.""" + + +@pytest.mark.asyncio +async def test_describe_types_and_members(server): + derived = Derived3D(server) + derived.mcp_expose("thing") + + index = await ok(server, "object_types") + assert index["types"]["Derived3D"] == {"methods": ["spin"], "extends": "Base3D"} + assert index["types"]["Base3D"] == { + "about": "The base of 3D things.", + "properties": ["shape"], + "methods": ["reset"], + } + + full = await ok(server, "object_types", type_names=["Derived3D"]) + assert full["types"]["Derived3D"] == { + "extends": "Base3D", + "properties": {}, + "methods": { + "spin": { + "description": "Spin it.", + "parameters": { + "type": "object", + "properties": {"turns": {"type": "integer", "default": 1}}, + "additionalProperties": False, + }, + } + }, + } + inherited = await ok( + server, "object_types", type_names=["Derived3D"], include_inherited=True + ) + assert set(inherited["types"]["Derived3D"]["methods"]) == {"reset", "spin"} + assert inherited["types"]["Derived3D"]["properties"]["shape"] == { + "type": "string", + "read_only": True, + } + + member = await ok(server, "object_types", type_names=["Base3D.reset"]) + assert member == { + "types": { + "Base3D": { + "methods": { + "reset": { + "description": "Back to defaults.", + "parameters": { + "type": "object", + "properties": {}, + "additionalProperties": False, + }, + } + } + } + } + } + # Types only: an object's name is unknown here + message = await fails(server, "object_types", type_names=["thing"]) + assert "Unknown types or members: thing" in message + + # Every wrong name, in one reply + message = await fails(server, "object_types", type_names=["Nope", "Base3D.nope"]) + assert "Unknown types or members: Nope" in message + assert "No member Base3D.nope" in message + + +@pytest.mark.asyncio +async def test_describe_parameters_and_returns(server): + desk = Desk(server) + desk.mcp_expose("desk") + described = await ok(server, "object_types", type_names=["Desk"]) + desk = described["types"]["Desk"] + # Object-typed values read the same everywhere: {"object": "T"} + assert desk["properties"]["active"] == { + "anyOf": [{"object": "Viewport"}, {"type": "null"}], + "read_only": True, + } + pick = desk["methods"]["pick"] + assert pick["parameters"]["properties"]["view"] == {"object": "Viewport"} + assert pick["returns"] == {"object": "Viewport"} + + +def test_server_describe(server): + assert server.mcp.describe(Plain) is None + described = server.mcp.describe(Viewport) + assert described["properties"]["background"] == { + "type": "string", + "description": "Background color", + } + assert described["properties"]["size"]["read_only"] is True + assert described["methods"]["wait"]["parameters"]["properties"]["seconds"] == { + "type": "number", + "minimum": 0, + "default": 0, + } + + +@pytest.mark.asyncio +async def test_type_name_collision(server, caplog): + def make(): + class Shared(TrameComponent): + @mcp.method + def go(self): ... + + return Shared + + first, second = make(), make() + a, b = first(server), second(server) + a.mcp_expose("a") + b.mcp_expose("b") + with caplog.at_level(logging.WARNING, logger="trame_server.mcp"): + listed = (await ok(server, "object_list"))["objects"] + assert listed["a"]["type"] == "Shared" + qualified = listed["b"]["type"] + assert qualified.startswith("test_mcp_objects.") or qualified.endswith(".Shared") + assert qualified != "Shared" + assert "Two classes agents see are named Shared" in caplog.text + + +@pytest.mark.asyncio +async def test_type_name_collision_after_start(server, caplog): + server.mcp._start() + + def make(): + class Latecomer(TrameComponent): + @mcp.method + def go(self): ... + + return Latecomer + + first, second = make(), make() + late = first(server) + late.mcp_expose("late") + with caplog.at_level(logging.WARNING, logger="trame_server.mcp"): + await ok(server, "object_types") # warned once agents look at types + assert "Two classes agents see are named Latecomer" in caplog.text + assert first is not second + + +# ----------------------------------------------------------------------------- +# Class options +# ----------------------------------------------------------------------------- + + +class Holder(TrameComponent, mcp={"description": "Holds things.", "all_methods": True}): + def visible(self) -> str: + return "yes" + + def _private(self): ... + + @property + def count(self) -> int: + return 3 + + @mcp.method + def also(self): ... + + +class Wide(Holder, mcp={"methods": {"visible": False}, "properties": ["count"]}): + pass + + +@pytest.mark.asyncio +async def test_class_options(server): + built = [Holder(server), Wide(server)] # held: names hold them weakly + for i, obj in enumerate(built): + obj.mcp_expose(f"thing_{i}") + described = await ok(server, "object_types", type_names=["Holder", "Wide"]) + holder = described["types"]["Holder"] + assert holder["description"] == "Holds things." + assert set(holder["methods"]) == {"also", "visible"} + assert holder["properties"] == {} + # What a subclass drops of its base is false + wide = described["types"]["Wide"] + assert wide == { + "extends": "Holder", + "properties": {"count": {"type": "integer", "read_only": True}}, + "methods": {"visible": False}, + } + index = (await ok(server, "object_types"))["types"]["Wide"] + assert index == {"properties": ["count"], "lacks": ["visible"], "extends": "Holder"} + described = await ok( + server, "object_types", type_names=["Wide"], include_inherited=True + ) + assert set(described["types"]["Wide"]["methods"]) == {"also"} + assert len(built) == 2 + + +class Pipeline(TrameComponent): + @mcp.method(annotations={"destructiveHint": True}) + def clear(self) -> None: + """Remove every + node.""" + + @mcp.method(annotations={"readOnlyHint": True, "idempotentHint": True}) + def status(self) -> str: + return "ok" + + +@pytest.mark.asyncio +async def test_method_annotations_and_docstrings(server): + pipeline = Pipeline(server) + pipeline.mcp_expose("pipeline", read_only=True) + described = await ok(server, "object_types", type_names=["Pipeline"]) + methods = described["types"]["Pipeline"]["methods"] + assert list(methods) == ["clear", "status"] + assert methods["clear"]["description"] == "Remove every node." + assert methods["clear"]["annotations"] == {"destructiveHint": True} + # readOnlyHint is read_only + assert methods["status"]["read_only"] is True + assert methods["status"]["annotations"] == {"idempotentHint": True} + assert await ok(server, "object_call", target="pipeline", method="status") == "ok" + + +@pytest.mark.asyncio +async def test_mcp_describe_on_a_class_you_dont_own(server): + class Library(TrameComponent): + def reset(self) -> None: + """Reset it.""" + + def other(self): ... + + with pytest.raises(TypeError, match="declares nothing"): + Library(server).mcp_expose("lib") + Library.mcp_describe(description="From a library.", methods=["reset"]) + lib = Library(server) + lib.mcp_expose("lib") + value = await ok(server, "object_types", type_names=["Library"]) + assert value["types"]["Library"]["description"] == "From a library." + assert set(value["types"]["Library"]["methods"]) == {"reset"} + assert await ok(server, "object_call", target=lib._id, method="reset") is None + + +# ----------------------------------------------------------------------------- +# object_get, object_update, object_call +# ----------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_get_references_and_depth(server): + desk = Desk(server) + desk.mcp_expose("desk") + view = Viewport(server) + desk.views.append(view) # reached through desk.active + value = await ok(server, "object_get", target=desk._id) + # Not read: a reference with its type + assert read(value)["values"] == {"active": {"id": view._id, "type": "Viewport"}} + assert list(value["objects"]) == [desk._id] + value = await ok(server, "object_get", target=desk._id, depth=2) + # Read: a reference by id, the object in "objects" + assert read(value)["values"] == {"active": {"id": view._id}} + assert value["objects"][view._id]["values"]["background"] == "black" + value = await ok(server, "object_get", target=view._id, properties=["size"]) + assert read(value)["values"] == {"size": 300} + message = await fails(server, "object_get", target=view._id, properties=["nope"]) + assert "has no properties ['nope']" in message + + +@pytest.mark.asyncio +async def test_update(server): + view = Viewport(server) + view.mcp_expose("view") + assert await ok( + server, "object_update", target=view._id, values={"background": "white"} + ) == {"id": view._id, "updated": ["background"]} + assert view.background == "white" + + message = await fails( + server, "object_update", target=view._id, values={"size": 3, "nope": 1} + ) + assert message.startswith("Nothing was applied.") + assert "size (read-only)" in message + assert "nope (unknown)" in message + # Methods may change what can't be written + assert message.endswith( + "Writable: background. Methods: set_view_up, describe_view, wait, echo; " + 'object_types(["Viewport"]) says what they do.' + ) + message = await fails(server, "object_update", target=view._id, values={"nope": 1}) + assert message.endswith("Writable: background") + message = await fails( + server, "object_update", target=view._id, values={"background": 3} + ) + assert message == "Nothing was applied: background: expected string, got integer" + + message = await fails( + server, "object_update", target=view._id, values={"background": "boom"} + ) + assert message.startswith(f"Setting Viewport {view._id}.background failed") + + +@pytest.mark.asyncio +async def test_call(server): + view = Viewport(server) + view.mcp_expose("view") + assert ( + await ok( + server, + "object_call", + target=view._id, + method="set_view_up", + args={"axis": "z"}, + ) + == "up z" + ) + assert ( + await ok( + server, "object_call", target=view._id, method="wait", args={"seconds": 2} + ) + == "waited 2.0" + ) + message = await fails( + server, "object_call", target=view._id, method="set_view_up", args={"axis": "w"} + ) + assert message == ( + 'Invalid arguments for Viewport.set_view_up: args.axis: must be one of "x", ' + '"y", "z", got "w"' + ) + # A union keeps the value's own type: an int for float | int + for given, shown in ((20, "int 20"), (2.5, "float 2.5"), ("a", "str a")): + assert ( + await ok( + server, + "object_call", + target=view._id, + method="echo", + args={"value": given}, + ) + == shown + ) + message = await fails(server, "object_call", target=view._id, method="nope") + assert message == ( + f'{view._id} (type Viewport) has no method "nope"; ' + 'object_types(["Viewport"]) lists its methods.' + ) + + +@pytest.mark.asyncio +async def test_call_with_object_arguments(server): + desk = Desk(server) + desk.mcp_expose("desk") + first, second = Viewport(server), Viewport(server) + second.mcp_expose("second") + desk.views.extend([first, second]) + + value = await ok( + server, "object_call", target=desk._id, method="pick", args={"view": "second"} + ) + value = read(value) + assert value["id"] == second._id + assert value["names"] == ["second"] + assert desk.active is second + + message = await fails( + server, "object_call", target=desk._id, method="pick", args={"view": desk._id} + ) + assert message == f'args.view: "{desk._id}" has type Desk, not Viewport' + message = await fails( + server, "object_call", target=desk._id, method="pick", args={"view": "999"} + ) + assert "is not the id or name of an object agents can reach" in message + + +@pytest.mark.asyncio +async def test_first_class_tool_with_object_parameter(server): + right = Viewport(server, "white") + right.mcp_expose("right") + + @server.mcp.tool + def screenshot(viewer: Viewport) -> str: + """A picture of a viewer.""" + return f"picture of {viewer.background}" + + assert await ok(server, "screenshot", viewer="right") == "picture of white" + + @server.mcp.tool + def viewers() -> dict: + """The viewers.""" + return {"right": right, "count": 1} + + @server.mcp.tool + def first_viewer() -> Viewport: + """The first viewer.""" + return right + + value = await ok(server, "viewers") + assert value["result"] == {"right": {"id": right._id}, "count": 1} + assert value["objects"][right._id]["values"]["background"] == "white" + value = read(await ok(server, "first_viewer")) + assert value["values"]["background"] == "white" + assert value["names"] == ["right"] + tools = await server.mcp.protocol.handle( + {"jsonrpc": "2.0", "id": 2, "method": "tools/list"} + ) + schema = next( + t["inputSchema"] for t in tools["result"]["tools"] if t["name"] == "screenshot" + ) + # Clients read a plain string with what it names + assert schema["properties"]["viewer"] == { + "type": "string", + "description": "Viewport object, by id or name", + } + + +@pytest.mark.asyncio +async def test_a_concurrent_method_runs_while_a_call_holds_the_lock(server): + release = asyncio.Event() + log = [] + + class Worker(TrameComponent): + @mcp.method + async def hold(self) -> str: + log.append("hold:start") + await release.wait() + log.append("hold:end") + return "held" + + @mcp.method(read_only=True, concurrent=True) + async def peek(self) -> str: + log.append("peek") + return "peeked" + + worker = Worker(server) + worker.mcp_expose("worker") + client = server.mcp.client() + + def calling(method): + return client.call_tool("object_call", {"target": "worker", "method": method}) + + holding = asyncio.ensure_future(calling("hold")) + await asyncio.sleep(0) # hold takes the lock + result = await asyncio.wait_for(calling("peek"), 1) + assert result["content"][0]["text"] == "peeked" + release.set() + await holding + assert log == ["hold:start", "peek", "hold:end"] + + +# ----------------------------------------------------------------------------- +# Permissions +# ----------------------------------------------------------------------------- + + +class Secret(TrameComponent): + def __init__(self, server): + super().__init__(server) + self._path = "/a" + + @mcp.method(requires="files") + def load(self, path: str) -> str: + return path + + @mcp.property(requires="files") + def path(self) -> str: + return self._path + + @mcp.property + def script(self) -> str: + return "print()" + + @script.setter(requires="scripts") + def script(self, value): ... + + +@pytest.mark.asyncio +async def test_permissions(server): + permissions = server.mcp.permissions + permissions.declare("files", "Read files on this computer") + permissions.declare("scripts", "Run Python code the agent writes", requires="files") + secret = Secret(server) + secret.mcp_expose("secret") + + full = (await ok(server, "object_types", type_names=["Secret"]))["types"]["Secret"] + assert full["methods"]["load"]["requires"] == ["files"] + assert full["properties"]["path"]["requires"] == ["files"] + assert full["properties"]["script"]["requires"] == ["scripts"] + + permissions["scripts"] = False + full = (await ok(server, "object_types", type_names=["Secret"]))["types"]["Secret"] + assert full["properties"]["script"] == {"type": "string", "read_only": True} + message = await fails( + server, "object_update", target=secret._id, values={"script": "x"} + ) + assert "script (read-only: permission scripts is off)" in message + + permissions["files"] = False + full = (await ok(server, "object_types", type_names=["Secret"]))["types"]["Secret"] + assert "load" not in full["methods"] + assert "path" not in full["properties"] + value = read(await ok(server, "object_get", target=secret._id)) + assert value["values"] == {"script": "print()"} + message = await fails( + server, "object_call", target=secret._id, method="load", args={"path": "/"} + ) + assert message == ( + 'Secret.load: permission "files" (Read files on this computer) is off; only ' + "the people running the app can turn it on." + ) + assert "files (Read files on this computer)" in server.mcp._agent_instructions() + + +# ----------------------------------------------------------------------------- +# The state object +# ----------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_state(server): + state = server.state + state.resolution = 6 + state.color = "red" + state.ready() + seen = [] + + @state.change("resolution") + def on_resolution(resolution, **_): + seen.append(resolution) + + state.mcp_expose( + "resolution", + type=int, + schema={"minimum": 3, "maximum": 60}, + description="Facets of the cone", + ) + state.mcp_expose("color", read_only=True) + + assert await ok(server, "object_list") == {"objects": {"state": {"type": "State"}}} + index = await ok(server, "object_types") + assert index["types"]["State"] == {"properties": ["resolution", "color"]} + full = await ok(server, "object_types", type_names=["State"]) + assert full["types"]["State"]["properties"] == { + "resolution": { + "type": "integer", + "minimum": 3, + "maximum": 60, + "description": "Facets of the cone", + }, + "color": {"read_only": True}, + } + + assert await ok(server, "object_get", target="state") == { + "result": {"id": "state"}, + "objects": { + "state": {"type": "State", "values": {"resolution": 6, "color": "red"}} + }, + } + assert await ok( + server, "object_update", target="state", values={"resolution": 24} + ) == {"id": "state", "updated": ["resolution"]} + assert state.resolution == 24 + assert seen == [24] + + message = await fails( + server, "object_update", target="state", values={"resolution": 2, "color": "b"} + ) + assert "color (read-only)" in message + message = await fails( + server, "object_update", target="state", values={"resolution": 2} + ) + assert message == "Nothing was applied: resolution: must be at least 3, got 2" + assert state.resolution == 24 + + state.mcp_unexpose("resolution", "color") + assert "Unknown object" in await fails(server, "object_get", target="state") + + +@pytest.mark.asyncio +async def test_state_of_child_servers(server): + child = server.create_child_server(prefix="left_") + child.state.background = "black" + child.state.mcp_expose("background", type=str) + value = read(await ok(server, "object_get", target="state")) + assert value["values"] == {"left_background": "black"} + await ok(server, "object_update", target="state", values={"left_background": "w"}) + assert child.state.background == "w" + + +@pytest.mark.asyncio +async def test_child_servers_translate_names(server): + """As they do state keys: a component built on each child server exposes + each instance under its own name.""" + left = server.create_child_server(prefix="left_") + right = server.create_child_server(prefix="right_") + left_view, right_view = Viewport(left), Viewport(right) + left_view.mcp_expose("viewer") + right_view.mcp_expose("viewer") + left.mcp.expose_object(lambda: left_view, "active") + + assert (await ok(server, "object_list"))["objects"] == { + "left_viewer": {"type": "Viewport"}, + "right_viewer": {"type": "Viewport"}, + "left_active": {"type": "Viewport"}, + } + assert read(await ok(server, "object_get", target="right_viewer"))["id"] == ( + right_view._id + ) + with pytest.raises(ValueError, match="'left_viewer' already exposes"): + left.mcp.expose_object(right_view, "viewer") + + left.mcp.unexpose_object("active") + assert "left_active" not in (await ok(server, "object_list"))["objects"] + + +def test_context_cannot_be_exposed(server): + with pytest.raises(RuntimeError, match="server-only"): + server.context.mcp_expose("secret") + + +# ----------------------------------------------------------------------------- +# object_batch +# ----------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_batch(server): + desk = Desk(server) + view = Viewport(server) + desk.views.append(view) + desk.mcp_expose("desk") + server.state.mcp_expose("zoom", type=float) + + value = await ok( + server, + "object_batch", + steps=[ + {"call": "desk", "method": "pick", "args": {"view": view._id}, "as": "v"}, + {"update": "@v", "values": {"background": "white"}}, + {"call": "@v", "method": "set_view_up", "args": {"axis": "x"}}, + {"update": "state", "values": {"zoom": 2}}, + ], + ) + assert value == { + "results": [ + {"id": view._id}, + {"id": view._id, "updated": ["background"]}, + "up x", + {"id": "state", "updated": ["zoom"]}, + ], + "objects": {view._id: {"type": "Viewport"}}, # links only, and none + } + value = await ok( + server, + "object_batch", + steps=[{"call": "desk", "method": "pick", "args": {"view": view._id}}], + results="full", + ) + assert value["objects"][view._id]["values"]["background"] == "white" + assert view.background == "white" + assert server.state.zoom == 2.0 + + message = await fails( + server, + "object_batch", + steps=[ + {"call": "desk", "method": "pick", "args": {"view": view._id}}, + {"update": view._id, "values": {"background": "boom"}}, + ], + ) + assert message.startswith( + f"steps[1] (update {view._id}) failed: Setting Viewport {view._id}.background" + ) + assert ". Step 0 stays applied (no rollback); its result: " in message + assert ".." not in message + # The first step: its own error, alone + message = await fails( + server, + "object_batch", + steps=[{"update": view._id, "values": {"size": 3}}], + ) + assert message == ( + f"steps[0] (update {view._id}) failed: Nothing was applied. Rejected on " + f"Viewport {view._id}: size (read-only). Writable: background. Methods: " + f'set_view_up, describe_view, wait, echo; object_types(["Viewport"]) ' + f"says what they do." + ) + + message = await fails(server, "object_batch", steps=[{"new": "Viewport"}]) + assert 'no type "Viewport" can be built' in message + message = await fails( + server, "object_batch", steps=[{"update": "@x", "values": {}}] + ) + assert "names no earlier step" in message + + +class Tile(TrameComponent, mcp={"constructor": "default"}): + """A tile on the wall.""" + + @mcp.property + def label(self) -> str: + return getattr(self, "_label", "") + + @label.setter + def label(self, value: str): + self._label = value + + +class Poster(TrameComponent): + """A poster on the wall.""" + + @mcp.constructor + @classmethod + def hang(cls, server, title: str) -> "Poster": + """Hang a new poster.""" + poster = cls(server) + poster.title = title + return poster + + @mcp.method(read_only=True) + def caption(self) -> str: + return self.title + + +class Wall(TrameComponent): + def __init__(self, server): + super().__init__(server) + self._tiles = [] + self._posters = [] + + # Agents create the types they can reach, here through these properties + @mcp.property + def tiles(self) -> list[Tile]: + return list(self._tiles) + + @mcp.property + def posters(self) -> list[Poster]: + return list(self._posters) + + @mcp.method + def add(self, tile: Tile) -> int: + self._tiles.append(tile) + return len(self._tiles) + + +@pytest.mark.asyncio +async def test_agents_create_components(server): + wall = Wall(server) + wall.mcp_expose("wall") + + types = (await ok(server, "object_types", type_names=["Tile", "Poster"]))["types"] + tile_new, poster_new = types["Tile"]["constructor"], types["Poster"]["constructor"] + assert tile_new["parameters"]["properties"] == {} # Tile(server) + assert poster_new["description"] == "Hang a new poster." + assert list(poster_new["parameters"]["properties"]) == ["title"] + + value = await ok( + server, + "object_batch", + steps=[ + {"new": "Tile", "as": "t"}, + {"update": "@t", "values": {"label": "north"}}, + {"call": "wall", "method": "add", "args": {"tile": "@t"}}, + {"new": "Poster", "args": {"title": "Hello"}, "as": "p"}, + {"call": "@p", "method": "caption"}, + ], + ) + results = value["results"] + assert results[2] == 1 + assert results[4] == "Hello" + (tile,) = wall._tiles + assert isinstance(tile, Tile) + assert tile.server is server + assert tile.label == "north" + assert results[0] == {"id": tile._id} + + +# ----------------------------------------------------------------------------- +# Response budget and instructions +# ----------------------------------------------------------------------------- + + +class Gauge(TrameComponent): + def __init__(self, server, count): + super().__init__(server) + self.count = count + + @mcp.property + def samples(self) -> list[int]: + return list(range(self.count)) + + @mcp.property + def label(self) -> str: + return "x" * 1000 + + +@pytest.mark.asyncio +async def test_unrequested_values_are_capped(server): + gauge = Gauge(server, 240) + gauge.mcp_expose("gauge") + value = await ok(server, "object_get", target="gauge") + assert read(value)["values"]["samples"] == list(range(20)) + assert len(read(value)["values"]["label"]) == 400 + assert value["more"] == { + f"{gauge._id}.samples": '20 of 240 items shown; object_get("gauge", ' + 'properties=["samples"]) for all', + f"{gauge._id}.label": '400 of 1000 characters shown; object_get("gauge", ' + 'properties=["label"]) for all', + } + value = await ok(server, "object_get", target="gauge", properties=["samples"]) + assert read(value)["values"]["samples"] == list(range(240)) + assert "more" not in value + + +@pytest.mark.asyncio +async def test_replies_fit_the_budget(server): + server.mcp.max_response_tokens = 500 # 2000 characters + gauge = Gauge(server, 5000) + gauge.mcp_expose("gauge") + value = await ok( + server, "object_get", target=gauge._id, properties=["samples", "label"] + ) + assert len(json.dumps(value)) <= 2000 + assert "to fit the response budget" in value["more"][f"{gauge._id}.samples"] + with pytest.raises(ValueError, match="at least 500"): + server.mcp.max_response_tokens = 10 + + +class Ledger(TrameComponent): + @mcp.property + def entries(self) -> list[str]: + return ["y" * 300] * 30 + + +@pytest.mark.asyncio +async def test_cut_values_count_from_what_was_read(server): + server.mcp.max_response_tokens = 500 + ledger = Ledger(server) + ledger.mcp_expose("ledger") + value = await ok(server, "object_get", target=ledger._id) + # Capped to 20 items, then cut again to fit: the note counts the 30 + shown = len(read(value)["values"]["entries"]) + assert shown < 20 + assert value["more"][f"{ledger._id}.entries"] == ( + f'{shown} of 30 items shown; object_get("{ledger._id}", ' + f'properties=["entries"]) alone for more' + ) + + +@pytest.mark.asyncio +async def test_instructions(server): + views = [Viewport(server) for _ in range(25)] # names hold them weakly + for i, view in enumerate(views): + view.mcp_expose(f"view_{i}") + response = await server.mcp.protocol.handle( + {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}} + ) + text = response["result"]["instructions"] + lines = text.split("\n") + assert lines[0] == ( + "A live trame app, shared with its users: they see your changes at once." + ) + assert lines[1].startswith("Objects: view_0 (Viewport), view_1 (Viewport)") + assert lines[1].endswith( + ", and 5 more; object_list() lists them, object_types() describes their types." + ) + assert len(text) < 1000 + + +@pytest.mark.asyncio +async def test_object_list_pages(server): + server.mcp.max_response_tokens = 500 # 2000 characters + views = [Viewport(server) for _ in range(60)] # names hold them weakly + for i, view in enumerate(views): + view.mcp_expose(f"view_{i}", description="One of many views.") + page = await ok(server, "object_list") + assert len(json.dumps(page, separators=(",", ":"))) <= 2000 + shown = len(page["objects"]) + assert 1 < shown < 60 + assert page["more"] == { + "objects": ( + f"1-{shown} of 60 listed to fit the response budget; " + f"object_list(offset={shown}) lists the next ones" + ) + } + + listed = dict(page["objects"]) + while "more" in page: + page = await ok(server, "object_list", offset=len(listed)) + listed.update(page["objects"]) + assert list(listed) == [f"view_{i}" for i in range(60)] + assert await ok(server, "object_list", offset=60) == {"objects": {}} + + +# ----------------------------------------------------------------------------- +# First-class tools of a component built once +# ----------------------------------------------------------------------------- + + +class Console(TrameComponent): + """Only first-class tools: not an object.""" + + @mcp.tool(read_only=True) + def status(self) -> dict: + """How the app is doing.""" + return {"ok": True} + + @mcp.resource("app://guide", mime_type="text/markdown") + def guide(self) -> str: + """How to use the app.""" + return "# Guide" + + @mcp.prompt + def review(self, focus: str = "all") -> str: + """Review the session.""" + return f"Review {focus}" + + +class Studio(TrameComponent): + @mcp.method + @mcp.tool(name="studio_reset") + def reset(self) -> str: + """Back to defaults.""" + return "reset" + + +async def rpc(server, method, **params): + response = await server.mcp.protocol.handle( + {"jsonrpc": "2.0", "id": 1, "method": method, "params": params} + ) + return response["result"] + + +@pytest.mark.asyncio +async def test_first_class_markers(server): + console = Console(server) + tools = (await rpc(server, "tools/list"))["tools"] + assert [t["name"] for t in tools] == ["status"] # not an object: no object tools + assert tools[0]["annotations"] == {"readOnlyHint": True} + assert await ok(server, "status") == {"ok": True} + resources = (await rpc(server, "resources/list"))["resources"] + assert resources == [ + { + "uri": "app://guide", + "name": "guide", + "description": "How to use the app.", + "mimeType": "text/markdown", + } + ] + read = await rpc(server, "resources/read", uri="app://guide") + assert read["contents"][0]["text"] == "# Guide" + prompt = await rpc(server, "prompts/get", name="review", arguments={"focus": "x"}) + assert prompt["messages"][0]["content"]["text"] == "Review x" + + # Built once: a second one would take the same names + with pytest.raises(ValueError, match="for components built once"): + Console(server) + assert server.mcp._tools["status"].fn.__self__ is console + + # Unbinding removes them, so another can be built + console._unbind_annotated_methods() + assert "status" not in server.mcp._tools + assert not server.mcp._resources + assert not server.mcp._prompts + second = Console(server) + assert server.mcp._tools["status"].fn.__self__ is second + + +@pytest.mark.asyncio +async def test_a_method_can_be_both(server): + studio = Studio(server) + assert await ok(server, "studio_reset") == "reset" # a tool needs no name + studio.mcp_expose("studio") + assert await ok(server, "object_call", target=studio._id, method="reset") == "reset" + + +@pytest.mark.asyncio +async def test_object_values_are_ids(server): + desk = Desk(server) + desk.mcp_expose("desk") + message = await fails( + server, + "object_call", + target=desk._id, + method="pick", + args={"view": {"id": "3"}}, + ) + assert message == ( + "Invalid arguments for Desk.pick: args.view: expected the id of an object " + "of type Viewport, got object" + ) diff --git a/tests/test_mcp_permissions.py b/tests/test_mcp_permissions.py new file mode 100644 index 0000000..a101dee --- /dev/null +++ b/tests/test_mcp_permissions.py @@ -0,0 +1,340 @@ +"""Permissions (``mcp.permissions``, ``requires=``, ``--mcp-deny``): powers +the people running an application withhold from agents.""" + +import itertools +import logging + +import pytest +from trame.app import get_server + +from trame_server.mcp import PermissionDeniedError +from trame_server.mcp.permissions import Permissions, launch_names + +_counter = itertools.count() + +LAUNCH_VARIABLES = ( + "TRAME_ARGS", + "TRAME_MCP_HTTP_TOKEN", + "TRAME_MCP_LOG", + "TRAME_MCP_DENY", +) + + +def rpc(method, params=None): + return {"jsonrpc": "2.0", "id": 1, "method": method, "params": params or {}} + + +def off(name, description): + return ( + f'{name}: permission "{description[0]}" ({description[1]}) is off; ' + f"only the people running the app can turn it on." + ) + + +@pytest.fixture(autouse=True) +def _no_launch_options(monkeypatch): + for name in LAUNCH_VARIABLES: + monkeypatch.delenv(name, raising=False) + + +@pytest.fixture +def mcp(): + mcp = get_server(f"test_mcp_permissions_{next(_counter)}").mcp + mcp.enabled = True + return mcp + + +@pytest.fixture +def permissions(mcp): + mcp.permissions.declare("files", "Read files and folders") + mcp.permissions.declare("scripts", "Run code the agent writes", requires="files") + return mcp.permissions + + +# ----------------------------------------------------------------------------- +# Declaring and setting +# ----------------------------------------------------------------------------- + + +def test_a_permission_is_allowed_only_with_what_it_requires(permissions): + assert dict(permissions) == {"files": True, "scripts": True} + permissions["files"] = False # a script could read files anyway + assert permissions.denied == ("files", "scripts") + refused = r"denied permission \('scripts' requires 'files'\); allow them" + with pytest.raises(ValueError, match=refused): + permissions["scripts"] = True # never brings files back by itself + assert permissions.denied == ("files", "scripts") + permissions.update(files=True, scripts=True) + assert permissions.denied == () + permissions["scripts"] = False + assert permissions.denied == ("scripts",) + + with pytest.raises(ValueError, match=refused): + permissions.update(scripts=True, files=False) + assert permissions.denied == ("scripts",) # all or none + with pytest.raises(KeyError, match="No permission 'script' is declared"): + permissions["script"] = False + with pytest.raises(KeyError, match="No permission 'script' is declared"): + permissions["script"] + + permissions.update({"scripts": True}, files=True) + assert permissions.denied == () + + +def test_declarations_are_checked(mcp): + permissions = mcp.permissions + with pytest.raises(ValueError, match="not declared"): + permissions.declare("scripts", "Run code", requires="files") + permissions.declare("files", "Read files", allowed=False) + permissions.declare("scripts", "Run code", requires=["files"]) + assert permissions["scripts"] is False # never allowed without what it requires + with pytest.raises(ValueError, match="already declared"): + permissions.declare("files", "Read files") + with pytest.raises(TypeError, match="non-empty string"): + permissions.declare("", "Nothing") + with pytest.raises(TypeError, match="requires must be a permission name"): + permissions.declare("net", "Network", requires=3) + + def ping() -> str: + return "pong" + + with pytest.raises(TypeError, match="requires must be a permission name"): + mcp.tool(requires=3)(ping) + with pytest.raises(TypeError, match="requires must be a permission name"): + mcp.tool(ping, requires=[""]) + assert mcp.get_tool("ping") is None + + +def test_what_is_missing_is_named_through_the_chain(): + permissions = Permissions() + permissions.declare("net", "Use the network", allowed=False) + permissions.declare("cloud", "Use the cloud", requires="net") + permissions.declare("backup", "Back up to the cloud", requires="cloud") + with pytest.raises(ValueError, match="'backup' requires 'net', 'cloud'"): + permissions["backup"] = True + permissions.update(net=True, cloud=True, backup=True) + assert permissions.denied == () + permissions["net"] = False + assert permissions.denied == ("net", "cloud", "backup") + + +def test_a_mapping_of_what_is_allowed(): + permissions = Permissions() + assert len(permissions) == 0 + permissions.declare("files", "Read files") + permissions.declare("net", "Use the network", allowed=False) + assert list(permissions) == ["files", "net"] + assert repr(permissions) == "Permissions({'files': True, 'net': False})" + assert permissions.description("net") == "Use the network" + assert permissions.allows([]) + assert permissions.allows(["files"]) + assert not permissions.allows(["files", "net"]) + assert not permissions.allows(["undeclared"]) + + +# ----------------------------------------------------------------------------- +# Tools +# ----------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_denied_tool_is_not_offered_and_refused(mcp, permissions): + @mcp.tool(requires="files") + def list_files(path: str) -> list: + return [path] + + @mcp.tool + def ping() -> str: + return "pong" + + permissions["files"] = False + listed = await mcp.protocol.handle(rpc("tools/list")) + assert [t["name"] for t in listed["result"]["tools"]] == ["ping"] + assert mcp.list_tools() == [mcp.get_tool("ping")] + + # A client that listed it earlier still calls it + reply = await mcp.protocol.handle( + rpc("tools/call", {"name": "list_files", "arguments": {"path": "/"}}) + ) + assert reply["result"]["isError"] + assert reply["result"]["content"][0]["text"] == off( + "list_files", ("files", "Read files and folders") + ) + + permissions["files"] = True + assert await mcp.protocol._execute(mcp.get_tool("list_files"), {"path": "/"}) == [ + "/" + ] + + +@pytest.mark.asyncio +async def test_a_tool_needing_several_permissions(mcp, permissions): + permissions.declare("net", "Use the network") + + @mcp.tool(requires=["files", "net"]) + def upload(path: str) -> str: + return path + + assert mcp.get_tool("upload").requires == ("files", "net") + permissions["net"] = False + assert mcp.list_tools() == [] + with pytest.raises(PermissionDeniedError) as refused: + await mcp.protocol._execute(mcp.get_tool("upload"), {"path": "/"}) + assert str(refused.value) == off("upload", ("net", "Use the network")) + + +@pytest.mark.asyncio +async def test_a_call_checks_what_its_arguments_need(mcp, permissions): + @mcp.tool + def run(external: bool = False) -> str: + if external: + mcp.permissions.require("files", what="run(external)") + return "ran" + + permissions["files"] = False + assert await mcp.protocol._execute(mcp.get_tool("run"), {}) == "ran" + with pytest.raises(PermissionDeniedError) as refused: + await mcp.protocol._execute(mcp.get_tool("run"), {"external": True}) + assert str(refused.value) == off( + "run(external)", ("files", "Read files and folders") + ) + + # The agent reads the message as is, without the exception's type + reply = await mcp.protocol.handle( + rpc("tools/call", {"name": "run", "arguments": {"external": True}}) + ) + assert reply["result"]["content"][0]["text"] == str(refused.value) + + +def test_require_names_the_refused_action(permissions): + permissions.require() # nothing needed + permissions.require("files", "scripts") + permissions["scripts"] = False + with pytest.raises( + PermissionDeniedError, match=r'^This call: permission "scripts"' + ): + permissions.require("files", "scripts") + with pytest.raises( + PermissionDeniedError, + match=r'^Export: permission "printer" was never declared by the app$', + ): + permissions.require("printer", what="Export") + + +@pytest.mark.asyncio +async def test_a_tool_needing_an_undeclared_permission_is_never_offered(mcp, caplog): + @mcp.tool(requires="network") + def fetch() -> str: + return "fetched" + + assert mcp.list_tools() == [] + with caplog.at_level(logging.WARNING, logger="trame_server.mcp"): + mcp._start() + assert ( + "Tool fetch requires undeclared permissions ['network']; it is never offered" + in caplog.text + ) + reply = await mcp.protocol.handle(rpc("tools/call", {"name": "fetch"})) + assert reply["result"]["isError"] + assert reply["result"]["content"][0]["text"] == ( + 'fetch: permission "network" was never declared by the app' + ) + + +@pytest.mark.asyncio +async def test_the_instructions_say_what_is_turned_off(mcp, permissions): + async def instructions(): + reply = await mcp.protocol.handle(rpc("initialize")) + return reply["result"]["instructions"] + + assert "Turned off" not in await instructions() + permissions["scripts"] = False + assert ( + "Turned off in this application: scripts (Run code the agent writes). What " + "needs them is not offered, and calls to it are refused." + ) in await instructions() + + +# ----------------------------------------------------------------------------- +# Launch +# ----------------------------------------------------------------------------- + + +def _launched(monkeypatch, argv, env=None): + # trame reads TRAME_ARGS instead of sys.argv when running under pytest + monkeypatch.setenv("TRAME_ARGS", argv) + if env is not None: + monkeypatch.setenv("TRAME_MCP_DENY", env) + mcp = get_server(f"test_mcp_permissions_launch_{next(_counter)}").mcp + mcp.enabled = True + mcp.permissions.declare("files", "Read files") + mcp.permissions.declare("scripts", "Run code", requires="files") + return mcp + + +def test_launch_names(): + assert launch_names(["a, b", None, "", " c ", ",,d"]) == ("a", "b", "c", "d") + assert launch_names([]) == () + + +def test_the_command_line_denies_for_the_whole_run(monkeypatch): + mcp = _launched(monkeypatch, "--mcp-deny files") + permissions = mcp.permissions + assert permissions.denied == permissions.locked == ("files", "scripts") + permissions["scripts"] = True + assert permissions.denied == ("files", "scripts") + mcp._start() + permissions.update(files=True, scripts=True) # a UI switch, once running + assert permissions.denied == ("files", "scripts") + + # Denying scripts leaves files to the application + mcp = _launched(monkeypatch, "--mcp-deny scripts") + permissions = mcp.permissions + assert permissions.locked == ("scripts",) + mcp._start() + permissions["files"] = False + permissions["scripts"] = True # locked: nothing changes + assert permissions.denied == ("files", "scripts") + permissions["files"] = True + assert permissions.denied == ("scripts",) + + +def test_a_permission_declared_later_is_denied_too(monkeypatch): + mcp = _launched(monkeypatch, "--mcp-deny net") + mcp.permissions.declare("net", "Use the network") + mcp.permissions.declare("cloud", "Use the cloud", requires="net") + assert mcp.permissions.denied == ("net", "cloud") + assert mcp.permissions.locked == ("net", "cloud") + + +def test_the_environment_denies_too(monkeypatch): + mcp = _launched( + monkeypatch, "--mcp-deny files --mcp-deny net,cloud", env="scripts, " + ) + assert mcp.permissions._launch_denied == {"files", "net", "cloud", "scripts"} + + +@pytest.mark.asyncio +async def test_an_unknown_name_at_launch_refuses_every_call(monkeypatch, caplog): + mcp = _launched(monkeypatch, "--mcp-deny script") + + @mcp.tool + def ping() -> str: + return "pong" + + # Until the server starts, the application may still declare it + assert await mcp.protocol._execute(mcp.get_tool("ping"), {}) == "pong" + + with caplog.at_level(logging.ERROR, logger="trame_server.mcp"): + mcp._start() + assert ( + "--mcp-deny names script, which this application does not declare " + "(it declares: files, scripts); every agent call is refused" + ) in caplog.text + with pytest.raises(PermissionDeniedError, match="launched denying script"): + await mcp.protocol._execute(mcp.get_tool("ping"), {}) + reply = await mcp.protocol.handle(rpc("tools/call", {"name": "ping"})) + assert reply["result"]["content"][0]["text"] == ( + "Agent access is refused: the application was launched denying script, " + "which it does not declare" + ) diff --git a/tests/test_mcp_progress.py b/tests/test_mcp_progress.py new file mode 100644 index 0000000..25c2a7f --- /dev/null +++ b/tests/test_mcp_progress.py @@ -0,0 +1,378 @@ +"""Progress (``Progress`` parameters): what a long tool or method reports, +sent to the agent that asked for it.""" + +import asyncio +import itertools +import json +import logging + +import aiohttp +import pytest +from trame.app import get_server +from trame_common.decorators import mcp as mcp_decorators +from trame_common.obj.component import TrameComponent + +from trame_server.mcp import Progress +from trame_server.mcp.progress import Reporter + +LAUNCH_VARIABLES = ( + "TRAME_ARGS", + "TRAME_MCP_HTTP_TOKEN", + "TRAME_MCP_LOG", + "TRAME_MCP_DENY", +) + +_counter = itertools.count() + + +def new_server(): + return get_server(f"test_mcp_progress_{next(_counter)}") + + +@pytest.fixture(autouse=True) +def _no_launch_options(monkeypatch): + for name in LAUNCH_VARIABLES: + monkeypatch.delenv(name, raising=False) + + +class Loader(TrameComponent): + """Loads files.""" + + @mcp_decorators.method + async def load(self, path: str, progress: Progress | None = None) -> str: + """Load a file.""" + self.received = progress + if progress: + await progress(1, total=2, message="Reading") + await asyncio.sleep(0.15) # past the interval: sent apart + await progress(2, total=2, message="Done") + return f"Loaded {path}" + + @mcp_decorators.method + async def scan(self, progress: Progress | None = None) -> str: + """Find files, not knowing how many there are.""" + if progress: + for found in (1, 2): + await progress(found, message=f"Found {found}") + await asyncio.sleep(0.15) + return "Scanned" + + +@pytest.fixture +def mcp(): + server = new_server() + mcp = server.mcp + mcp.configure(enabled=True, transports=[]) + + @mcp.tool + async def crunch(steps: int, progress: Progress | None = None) -> str: + """Crunch some numbers.""" + for step in range(1, steps + 1): + await progress(step, total=steps, message=f"Step {step}") + await asyncio.sleep(0.15) + return "Crunched" + + loader = Loader(server) + loader.mcp_expose("loader") + mcp.loader = loader # kept alive for the test + return mcp + + +def collector(): + reports = [] + + def on_progress(value, total, message): + reports.append((value, total, message)) + + return reports, on_progress + + +# ----------------------------------------------------------------------------- +# Declaring +# ----------------------------------------------------------------------------- + + +def test_agents_dont_see_the_parameter(mcp): + (tool,) = [t for t in mcp.list_tools() if t.name == "crunch"] + schema = mcp._input_schema(tool) + assert list(schema["properties"]) == ["steps"] + assert schema["additionalProperties"] is False + + parameters = mcp.describe(Loader)["methods"]["load"]["parameters"] + assert list(parameters["properties"]) == ["path"] + + +def test_reports_are_awaited_so_it_takes_async(mcp): + def sync_crunch(progress: Progress | None = None) -> str: # noqa: ARG001 + return "no" + + with pytest.raises(TypeError, match="Tool 'sync_crunch': a Progress parameter"): + mcp.tool(sync_crunch) + + class SyncLoader(TrameComponent): + @mcp_decorators.method + def load(self, progress: Progress | None = None) -> str: # noqa: ARG002 + return "no" + + with pytest.raises(TypeError, match=r"SyncLoader\.load: a Progress parameter"): + mcp.describe(SyncLoader) + + +@pytest.mark.asyncio +async def test_the_apps_own_calls_get_the_default(mcp): + assert await mcp.loader.load("a.vtk") == "Loaded a.vtk" + assert mcp.loader.received is None + + +# ----------------------------------------------------------------------------- +# Reporting +# ----------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_tool_reports_to_the_agent_that_asked(mcp): + reports, on_progress = collector() + result = await mcp.client().call_tool( + "crunch", {"steps": 3}, on_progress=on_progress + ) + assert result["content"][0]["text"] == "Crunched" + assert reports == [(1, 3, "Step 1"), (2, 3, "Step 2"), (3, 3, "Step 3")] + + +@pytest.mark.asyncio +async def test_without_asking_nothing_is_reported(mcp): + client = mcp.client() + result = await client.call_tool("crunch", {"steps": 2}) + assert result["content"][0]["text"] == "Crunched" # a reporter that does nothing + + await client.call_tool( + "object_call", {"target": "loader", "method": "load", "args": {"path": "a"}} + ) + assert type(mcp.loader.received) is Progress # not None: an agent's call + + +@pytest.mark.asyncio +async def test_a_method_reports_through_object_call(mcp): + reports, on_progress = collector() + result = await mcp.client().call_tool( + "object_call", + {"target": "loader", "method": "load", "args": {"path": "a.vtk"}}, + on_progress=on_progress, + ) + assert result["content"][0]["text"] == "Loaded a.vtk" + assert reports == [(1, 2, "Reading"), (2, 2, "Done")] + + +@pytest.mark.asyncio +async def test_a_batch_reports_its_steps(mcp): + """Step k of n, a method's own report counting within its step.""" + reports, on_progress = collector() + steps = [ + {"call": "loader", "method": "load", "args": {"path": "a"}}, + {"call": "loader", "method": "load", "args": {"path": "b"}}, + ] + result = await mcp.client().call_tool( + "object_batch", {"steps": steps}, on_progress=on_progress + ) + assert "isError" not in result + values = [value for value, _, _ in reports] + assert values == sorted(set(values)) # always growing + assert all(total == 2 for _, total, _ in reports) + assert values[-1] == 2 + assert (0.5, 2, "Reading") in reports # the first load, half way + + +@pytest.mark.asyncio +async def test_a_batch_step_without_a_total_still_reports(mcp): + """Each report moves the value a little within its step, so it goes out + while the step runs, with its message.""" + reports, on_progress = collector() + steps = [ + {"call": "loader", "method": "load", "args": {"path": "a"}}, + {"call": "loader", "method": "scan"}, + ] + result = await mcp.client().call_tool( + "object_batch", {"steps": steps}, on_progress=on_progress + ) + assert "isError" not in result + assert reports[-3:] == [ + (1.00000001, 2, "Found 1"), + (1.00000002, 2, "Found 2"), + (2, 2, None), + ] + + +@pytest.mark.asyncio +async def test_a_failing_callback_doesnt_stop_the_call(mcp, caplog): + def on_progress(*_): + msg = "broken UI" + raise RuntimeError(msg) + + with caplog.at_level(logging.ERROR, logger="trame_server.mcp"): + result = await mcp.client().call_tool( + "crunch", {"steps": 1}, on_progress=on_progress + ) + assert result["content"][0]["text"] == "Crunched" + assert "on_progress failed" in caplog.text + + +@pytest.mark.asyncio +async def test_an_async_callback(mcp): + reports = [] + + async def on_progress(value, _total, _message): + await asyncio.sleep(0) + reports.append(value) + + await mcp.client().call_tool("crunch", {"steps": 2}, on_progress=on_progress) + assert reports == [1, 2] + + +# ----------------------------------------------------------------------------- +# The reporter's rules +# ----------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_reports_are_combined_the_latest_winning(): + sent = [] + + async def send(fields): + sent.append(fields) + + reporter = Reporter(send, interval=0.05) + for value in range(1, 11): + await reporter(value, total=10) + assert sent == [{"progress": 1, "total": 10}] # the rest is due later + await asyncio.sleep(0.08) + assert sent[-1] == {"progress": 10, "total": 10} + assert len(sent) == 2 + + +@pytest.mark.asyncio +async def test_a_value_must_grow(): + sent = [] + + async def send(fields): + sent.append(fields) + + reporter = Reporter(send, interval=0) + await reporter(5) + await reporter(5, message="Still at five") # dropped, its message kept + await reporter(3) + await reporter(6) + assert sent == [{"progress": 5}, {"progress": 6, "message": "Still at five"}] + + +@pytest.mark.asyncio +async def test_what_is_pending_goes_before_the_result_then_nothing(): + sent = [] + + async def send(fields): + sent.append(fields) + + reporter = Reporter(send, interval=0.05) + await reporter(1) + await reporter(2) # due later + await reporter.close() # the result is ready + assert sent == [{"progress": 1}, {"progress": 2}] + await reporter(3) + await asyncio.sleep(0.08) + assert sent == [{"progress": 1}, {"progress": 2}] + + +@pytest.mark.asyncio +async def test_a_send_that_fails_stops_the_reports_not_the_call(): + calls = [] + + async def send(fields): + calls.append(fields) + raise ConnectionResetError + + reporter = Reporter(send, interval=0) + await reporter(1) + await reporter(2) + assert calls == [{"progress": 1}] + + +# ----------------------------------------------------------------------------- +# Over HTTP: Server-Sent Events once there is something to report +# ----------------------------------------------------------------------------- + + +async def start(server): + server.start( + exec_mode="task", + port=0, + open_browser=False, + show_connection_info=False, + disable_logging=True, + ) + await server.ready + + +def tool_call(name, arguments, token=None): + meta = { + "io.modelcontextprotocol/protocolVersion": "2026-07-28", + "io.modelcontextprotocol/clientCapabilities": {}, + } + if token is not None: + meta["progressToken"] = token + message = { + "jsonrpc": "2.0", + "id": 7, + "method": "tools/call", + "params": {"name": name, "arguments": arguments, "_meta": meta}, + } + headers = { + "MCP-Protocol-Version": "2026-07-28", + "Mcp-Method": "tools/call", + "Mcp-Name": name, + "Accept": "application/json, text/event-stream", + } + return message, headers + + +def events(text): + return [ + json.loads(line[len("data: ") :]) + for line in text.splitlines() + if line.startswith("data: ") + ] + + +@pytest.mark.asyncio +async def test_over_http(mcp): + server = mcp.server + mcp.transports = ["http"] + await start(server) + url = f"http://localhost:{server.port}/mcp" + try: + async with aiohttp.ClientSession() as session: + # Asked for: the notifications, then the response, as events + message, headers = tool_call("crunch", {"steps": 2}, token="t1") + async with session.post(url, json=message, headers=headers) as resp: + assert resp.status == 200 + assert resp.headers["Content-Type"].startswith("text/event-stream") + assert resp.headers["X-Accel-Buffering"] == "no" + received = events(await resp.text()) + *notifications, response = received + assert [n["params"] for n in notifications] == [ + {"progressToken": "t1", "progress": 1, "total": 2, "message": "Step 1"}, + {"progressToken": "t1", "progress": 2, "total": 2, "message": "Step 2"}, + ] + assert {n["method"] for n in notifications} == {"notifications/progress"} + assert response["id"] == 7 + assert response["result"]["content"][0]["text"] == "Crunched" + + # Not asked for, or a client that takes no event stream: plain JSON + message, headers = tool_call("crunch", {"steps": 1}) + async with session.post(url, json=message, headers=headers) as resp: + assert resp.headers["Content-Type"].startswith("application/json") + message, headers = tool_call("crunch", {"steps": 1}, token="t2") + headers["Accept"] = "application/json" + async with session.post(url, json=message, headers=headers) as resp: + assert resp.headers["Content-Type"].startswith("application/json") + assert (await resp.json())["result"]["content"][0]["text"] == "Crunched" + finally: + await server.stop() diff --git a/tests/test_mcp_protocol.py b/tests/test_mcp_protocol.py new file mode 100644 index 0000000..55d49ba --- /dev/null +++ b/tests/test_mcp_protocol.py @@ -0,0 +1,1037 @@ +"""MCP JSON-RPC dispatch (``mcp.protocol``), without HTTP.""" + +import asyncio +import base64 +import itertools +import json +from typing import Annotated, Literal + +import pytest +from trame.app import get_server + +from trame_server import __version__ +from trame_server.mcp import Image, MCPDisabledError, ToolError +from trame_server.mcp.protocol import ( + INTERNAL_ERROR, + INVALID_PARAMS, + INVALID_REQUEST, + LEGACY_VERSIONS, + METHOD_NOT_FOUND, + PROTOCOL_VERSIONS, + REQUEST_CANCELLED, + UNSUPPORTED_VERSION, +) + +LAUNCH_VARIABLES = ( + "TRAME_ARGS", + "TRAME_MCP_HTTP_TOKEN", + "TRAME_MCP_LOG", + "TRAME_MCP_DENY", +) + +_counter = itertools.count() + + +def rpc(method, params=None, rpc_id=1): + msg = {"jsonrpc": "2.0", "id": rpc_id, "method": method} + if params is not None: + msg["params"] = params + return msg + + +def modern(method, params=None, rpc_id=1, version="2026-07-28", capabilities=None): + """A request as clients of 2026-07-28 send it: its version and the + client's capabilities in _meta.""" + meta = { + "io.modelcontextprotocol/protocolVersion": version, + "io.modelcontextprotocol/clientInfo": {"name": "test-client", "version": "1"}, + "io.modelcontextprotocol/clientCapabilities": capabilities or {}, + } + return rpc(method, {**(params or {}), "_meta": meta}, rpc_id=rpc_id) + + +def call_tool(mcp, name, rpc_id=1, **arguments): + return mcp.protocol.handle( + rpc("tools/call", {"name": name, "arguments": arguments}, rpc_id=rpc_id) + ) + + +def text_of(response): + return response["result"]["content"][0]["text"] + + +def cancel_notification(request_id, reason=None): + params = {"requestId": request_id} + if reason is not None: + params["reason"] = reason + return {"jsonrpc": "2.0", "method": "notifications/cancelled", "params": params} + + +@pytest.fixture(autouse=True) +def _no_launch_options(monkeypatch): + for name in LAUNCH_VARIABLES: + monkeypatch.delenv(name, raising=False) + + +@pytest.fixture +def mcp(): + mcp = get_server(f"test_mcp_protocol_{next(_counter)}").mcp + mcp.enabled = True + return mcp + + +@pytest.fixture +def protocol(mcp): + @mcp.tool + def add(a: int, b: int = 1) -> int: + """Add two numbers.""" + return a + b + + @mcp.tool + async def async_echo(text: str) -> str: + return text + + @mcp.tool + def failing(): + msg = "boom" + raise RuntimeError(msg) + + @mcp.resource("trame://info", mime_type="application/json") + def info(): + return {"name": "test"} + + @mcp.prompt + def walkthrough(topic: Annotated[str, "Part of the app to explain"] = "app") -> str: + return f"Explain the {topic}" + + return mcp.protocol + + +# ----------------------------------------------------------------------------- +# Lifecycle +# ----------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_initialize(protocol, mcp): + response = await protocol.handle(rpc("initialize", {"capabilities": {}})) + result = response["result"] + assert result["protocolVersion"] == "2025-11-25" + assert set(result["capabilities"]) == {"tools", "resources", "prompts"} + assert result["serverInfo"] == {"name": mcp.server.name, "version": __version__} + + mcp.name = "viewer" + response = await protocol.handle(rpc("initialize")) + assert response["result"]["serverInfo"]["name"] == "viewer" + + +def test_protocol_versions(): + assert PROTOCOL_VERSIONS == ( + "2026-07-28", + "2025-11-25", + "2025-06-18", + "2025-03-26", + "2024-11-05", + ) + assert LEGACY_VERSIONS == PROTOCOL_VERSIONS[1:] # those opening with initialize + + +@pytest.mark.asyncio +@pytest.mark.parametrize("version", LEGACY_VERSIONS) +async def test_initialize_negotiates_older_versions(protocol, version): + response = await protocol.handle(rpc("initialize", {"protocolVersion": version})) + assert response["result"]["protocolVersion"] == version + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "version", + # 2026-07-28 has no handshake: initialize can't select it + ["1999-01-01", None, ["2025-06-18"], "2026-07-28"], +) +async def test_initialize_unknown_version_falls_back(protocol, version): + response = await protocol.handle(rpc("initialize", {"protocolVersion": version})) + assert response["result"]["protocolVersion"] == "2025-11-25" + + +@pytest.mark.asyncio +async def test_initialize_sends_app_then_generated_instructions(mcp): + async def instructions(): + response = await mcp.protocol.handle(rpc("initialize")) + return response["result"]["instructions"] + + # nothing exposed: only the live-app notice + text = await instructions() + assert ( + text + == "A live trame app, shared with its users: they see your changes at once." + ) + + mcp.instructions = "A library catalog.\n" + mcp.server.state.query = "" + mcp.server.state.mcp_expose("query", read_only=True) + lines = (await instructions()).split("\n") + assert lines[0] == "A library catalog." + assert lines[1] == "" + assert lines[2].startswith("A live trame app") + assert lines[3] == ( + "Objects: state (State); object_list() lists them, object_types() " + "describes their types." + ) + + +# ----------------------------------------------------------------------------- +# 2026-07-28: no handshake, every request names its version +# ----------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_discover(protocol, mcp): + mcp.instructions = "A test app." + response = await protocol.handle(modern("server/discover")) + result = response["result"] + assert result["resultType"] == "complete" + assert result["supportedVersions"] == list(PROTOCOL_VERSIONS) + assert set(result["capabilities"]) == {"tools", "resources", "prompts"} + assert result["instructions"].startswith("A test app.") + assert result["_meta"]["io.modelcontextprotocol/serverInfo"] == { + "name": mcp.server.name, + "version": __version__, + } + assert (result["ttlMs"], result["cacheScope"]) == (0, "private") + + +@pytest.mark.asyncio +async def test_modern_results(protocol, mcp): + """Every result says it is complete and who answered; lists and reads + also carry caching hints: none fresh, none shared.""" + server_info = {"name": mcp.server.name, "version": __version__} + listed = (await protocol.handle(modern("tools/list")))["result"] + assert listed["resultType"] == "complete" + assert listed["_meta"] == {"io.modelcontextprotocol/serverInfo": server_info} + assert (listed["ttlMs"], listed["cacheScope"]) == (0, "private") + assert "add" in [tool["name"] for tool in listed["tools"]] + + called = await protocol.handle( + modern("tools/call", {"name": "add", "arguments": {"a": 2}}) + ) + result = called["result"] + assert result["content"] == [{"type": "text", "text": "3"}] + assert result["resultType"] == "complete" + assert result["_meta"] == {"io.modelcontextprotocol/serverInfo": server_info} + assert "ttlMs" not in result + + # A tool's failure is a complete result too + failed = await protocol.handle(modern("tools/call", {"name": "failing"})) + assert failed["result"]["isError"] is True + assert failed["result"]["resultType"] == "complete" + + # Clients before 2026-07-28 get the results of their revision + legacy = (await protocol.handle(rpc("tools/list")))["result"] + assert not {"resultType", "ttlMs", "cacheScope", "_meta"} & set(legacy) + + +@pytest.mark.asyncio +async def test_an_unsupported_version(protocol): + response = await protocol.handle(modern("tools/list", version="1900-01-01")) + error = response["error"] + assert error["code"] == UNSUPPORTED_VERSION + assert error["data"] == { + "supported": list(PROTOCOL_VERSIONS), + "requested": "1900-01-01", + } + + # an earlier revision isn't spoken this way: it opens with initialize + response = await protocol.handle(modern("tools/list", version="2025-11-25")) + assert response["error"]["code"] == UNSUPPORTED_VERSION + assert "it opens with initialize" in response["error"]["message"] + + +@pytest.mark.asyncio +async def test_the_client_capabilities_are_required(protocol): + message = modern("tools/list") + del message["params"]["_meta"]["io.modelcontextprotocol/clientCapabilities"] + response = await protocol.handle(message) + assert response["error"]["code"] == INVALID_PARAMS + assert "clientCapabilities" in response["error"]["message"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["initialize", "ping"]) +async def test_the_handshake_and_ping_are_gone(protocol, method): + response = await protocol.handle(modern(method)) + assert response["error"]["code"] == METHOD_NOT_FOUND + assert protocol.handles(modern(method)) is False + assert protocol.handles(rpc(method)) is True # still served to older clients + assert protocol.handles(rpc("server/discover")) is False + + +@pytest.mark.asyncio +async def test_notifications_and_client_responses_get_no_response(protocol): + assert await protocol.handle(rpc("notifications/initialized", rpc_id=None)) is None + msg = {"jsonrpc": "2.0", "method": "notifications/initialized"} + assert await protocol.handle(msg) is None + # a reply to a request of the server's + assert await protocol.handle({"jsonrpc": "2.0", "id": 5, "result": {}}) is None + + +@pytest.mark.asyncio +async def test_malformed_requests(protocol): + for message in ({"id": 1, "method": "ping"}, {"jsonrpc": "1.0"}, [rpc("ping")]): + response = await protocol.handle(message) + assert response["id"] is None + assert response["error"]["code"] == INVALID_REQUEST + response = await protocol.handle(rpc("tools/list", params=["not", "a", "dict"])) + assert response["error"]["code"] == INVALID_PARAMS + response = await protocol.handle({"jsonrpc": "2.0", "id": 1, "method": 42}) + assert response["error"]["code"] == INVALID_REQUEST + response = await protocol.handle( + rpc("prompts/get", {"name": "walkthrough", "arguments": {"bogus": "x"}}) + ) + assert response["error"]["code"] == INVALID_PARAMS + + +@pytest.mark.asyncio +async def test_unknown_method(protocol): + response = await protocol.handle(rpc("bogus/method", rpc_id="a")) + assert response == { + "jsonrpc": "2.0", + "id": "a", + "error": { + "code": METHOD_NOT_FOUND, + "message": "Method not found: bogus/method", + }, + } + + +@pytest.mark.asyncio +async def test_ping(protocol): + response = await protocol.handle(rpc("ping")) + assert response == {"jsonrpc": "2.0", "id": 1, "result": {}} + + +# ----------------------------------------------------------------------------- +# Tools +# ----------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_tools_list(protocol, mcp): + @mcp.tool(read_only=True) + def peek() -> str: + return "ok" + + @mcp.tool(read_only=True, annotations={"openWorldHint": False}) + def look() -> str: + return "ok" + + response = await protocol.handle(rpc("tools/list")) + tools = {t["name"]: t for t in response["result"]["tools"]} + assert list(tools) == ["add", "async_echo", "failing", "peek", "look"] + assert tools["add"]["description"] == "Add two numbers." + assert tools["add"]["inputSchema"] == { + "type": "object", + "properties": { + "a": {"type": "integer"}, + "b": {"type": "integer", "default": 1}, + }, + "required": ["a"], + "additionalProperties": False, + } + assert "annotations" not in tools["add"] + assert tools["failing"]["description"] == "" + assert tools["peek"]["annotations"] == {"readOnlyHint": True} + assert tools["look"]["annotations"] == { + "openWorldHint": False, + "readOnlyHint": True, + } + + +@pytest.mark.asyncio +async def test_the_object_tools_are_listed_once_something_is_exposed(mcp): + response = await mcp.protocol.handle(rpc("tools/list")) + assert response["result"]["tools"] == [] + + mcp.server.state.size = 1 + mcp.server.state.mcp_expose("size", type=int) + response = await mcp.protocol.handle(rpc("tools/list")) + assert [t["name"] for t in response["result"]["tools"]] == [ + "object_list", + "object_types", + "object_get", + "object_update", + "object_call", + "object_batch", + ] + + +@pytest.mark.asyncio +async def test_tools_call_sync_and_async(protocol): + response = await protocol.handle( + rpc("tools/call", {"name": "add", "arguments": {"a": 2, "b": 3}}) + ) + assert response["result"] == {"content": [{"type": "text", "text": "5"}]} + + response = await protocol.handle( + rpc("tools/call", {"name": "async_echo", "arguments": {"text": "hi"}}) + ) + assert response["result"]["content"] == [{"type": "text", "text": "hi"}] + + +@pytest.mark.asyncio +async def test_tools_call_without_arguments(mcp): + @mcp.tool + def hello() -> str: + return "hi" + + response = await mcp.protocol.handle(rpc("tools/call", {"name": "hello"})) + assert text_of(response) == "hi" + + +@pytest.mark.asyncio +async def test_tools_call_errors_go_in_the_result(protocol, mcp): + @mcp.tool + def refuse(reason: str): + if reason == "tool": + msg = "Pick a smaller value; object_get shows the range." + raise ToolError(msg) + if reason == "value": + msg = "bad value" + raise ValueError(msg) + msg = "not yours" + raise PermissionError(msg) + + response = await protocol.handle(rpc("tools/call", {"name": "failing"})) + assert response["result"] == { + "content": [{"type": "text", "text": "RuntimeError: boom"}], + "isError": True, + } + # complete messages are sent as they are + response = await call_tool(mcp, "refuse", reason="tool") + assert text_of(response) == "Pick a smaller value; object_get shows the range." + response = await call_tool(mcp, "refuse", reason="permission") + assert text_of(response) == "not yours" + response = await call_tool(mcp, "refuse", reason="value") + assert text_of(response) == "ValueError: bad value" + + +@pytest.mark.asyncio +async def test_tools_call_bad_arguments(protocol, mcp): + @mcp.tool + def set_view_up(axis: Literal["x", "y", "z"]) -> str: + return axis + + async def call(name, arguments): + response = await protocol.handle( + rpc("tools/call", {"name": name, "arguments": arguments}) + ) + result = response["result"] + assert result["isError"] is True + return result["content"][0]["text"] + + assert "unexpected 'wrong'" in await call("add", {"a": 1, "wrong": 1}) + assert "missing required 'a'" in await call("add", {}) + assert await call("add", {"a": "1"}) == ( + "Invalid arguments for tool add: arguments.a: expected integer, got string" + ) + assert await call("add", [1]) == ( + "Invalid arguments for tool add: arguments: expected object, got array" + ) + assert await call("set_view_up", {"axis": "w"}) == ( + "Invalid arguments for tool set_view_up: arguments.axis: must be one of " + '"x", "y", "z", got "w"' + ) + + +@pytest.mark.asyncio +async def test_tools_call_type_error_in_body_is_not_blamed_on_arguments(mcp): + @mcp.tool + def measure(x: int) -> int: + return len(x) + + text = text_of(await call_tool(mcp, "measure", x=3)) + assert text.startswith("TypeError:") + assert "Invalid arguments" not in text + + +@pytest.mark.asyncio +async def test_tools_call_coerces_lossless_numbers(mcp): + @mcp.tool + def numbers(i: int, f: float) -> dict: + return {"i": type(i).__name__, "f": type(f).__name__} + + response = await call_tool(mcp, "numbers", i=2.0, f=2) + assert json.loads(text_of(response)) == {"i": "int", "f": "float"} + + +@pytest.mark.asyncio +async def test_tools_call_restores_declared_tuples(mcp): + @mcp.tool + def place(position: tuple[float, float, float]) -> str: + return type(position).__name__ + repr(position) + + response = await call_tool(mcp, "place", position=[1, 2, 3]) + assert text_of(response) == "tuple(1.0, 2.0, 3.0)" + + +@pytest.mark.asyncio +async def test_a_result_is_sent_once_unless_asked_twice(mcp): + @mcp.tool + def info() -> dict: + return {"name": "Å", "sizes": [1, 2]} + + response = await mcp.protocol.handle(rpc("tools/call", {"name": "info"})) + assert response["result"] == { + "content": [{"type": "text", "text": '{"name":"Å","sizes":[1,2]}'}] + } + mcp.structured_content = True + response = await mcp.protocol.handle(rpc("tools/call", {"name": "info"})) + assert response["result"]["structuredContent"] == {"name": "Å", "sizes": [1, 2]} + + +@pytest.mark.asyncio +async def test_results_become_content(mcp): + png = b"\x89PNG\r\n" + encoded = base64.b64encode(png).decode() + + @mcp.tool + def nothing() -> None: + pass + + @mcp.tool + def screenshot() -> Image: + return Image(png) + + @mcp.tool + def views() -> list: + return [Image(png, mime_type="image/jpeg"), "front"] + + @mcp.tool + def numbers() -> list: + return [1, 2.5, None] + + @mcp.tool + def mixed() -> dict: + return {"image": Image(png), "tags": {1}} + + async def result(name): + return (await call_tool(mcp, name))["result"] + + assert await result("nothing") == {"content": [{"type": "text", "text": "null"}]} + assert await result("screenshot") == { + "content": [{"type": "image", "data": encoded, "mimeType": "image/png"}] + } + assert await result("views") == { + "content": [ + {"type": "image", "data": encoded, "mimeType": "image/jpeg"}, + {"type": "text", "text": "front"}, + ] + } + assert await result("numbers") == { + "content": [{"type": "text", "text": "[1,2.5,null]"}] + } + assert (await result("mixed"))["content"][0]["text"] == ( + '{"image":"","tags":"{1}"}' + ) + + +@pytest.mark.asyncio +async def test_tools_call_unknown_tool(protocol): + response = await protocol.handle(rpc("tools/call", {"name": "nope"})) + assert response["error"] == { + "code": INVALID_PARAMS, + "message": "Unknown tool: nope", + } + + +@pytest.mark.asyncio +async def test_a_bound_method_is_a_tool(mcp): + class ProtoCounter: + def __init__(self): + self.count = 0 + + def increment(self, by: Annotated[int, "Step", {"minimum": 1}] = 1) -> int: + """Count up.""" + self.count += by + return self.count + + counter = ProtoCounter() + assert mcp.tool(counter.increment) == counter.increment + entry = mcp.get_tool("increment") + assert entry.description == "Count up." + response = await mcp.protocol.handle(rpc("tools/list")) + (descriptor,) = response["result"]["tools"] + assert descriptor["inputSchema"]["properties"] == { + "by": {"type": "integer", "description": "Step", "minimum": 1, "default": 1} + } + assert text_of(await call_tool(mcp, "increment", by=2)) == "2" + assert counter.count == 2 + assert "must be at least 1" in text_of(await call_tool(mcp, "increment", by=0)) + + +@pytest.mark.asyncio +async def test_the_object_tools_are_served_like_any_tool(mcp): + state = mcp.server.state + state.resolution = 6 + state.ready() # what starting the server does: flushes run from now on + state.mcp_expose("resolution", type=int, schema={"minimum": 3}) + seen = [] + + @state.change("resolution") + def on_resolution(resolution, **_): + seen.append(resolution) + + response = await call_tool( + mcp, "object_update", target="state", values={"resolution": 12} + ) + assert json.loads(text_of(response)) == {"id": "state", "updated": ["resolution"]} + assert state.resolution == 12 + assert seen == [12] + + response = await call_tool(mcp, "object_get", target="state") + reply = json.loads(text_of(response)) + assert reply["objects"]["state"]["values"] == {"resolution": 12} + + response = await call_tool( + mcp, "object_update", target="state", values={"resolution": 1} + ) + assert response["result"]["isError"] is True + assert state.resolution == 12 + + +# ----------------------------------------------------------------------------- +# Resources and prompts +# ----------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_resources(protocol, mcp): + @mcp.resource("trame://logo", name="logo", description="The logo") + def logo() -> bytes: + return b"\x00\x01" + + @mcp.resource("trame://readme") + async def readme() -> str: + """How to use it.""" + return "# Readme" + + @mcp.resource("trame://sizes") + def sizes() -> list: + return [1, 2] + + response = await protocol.handle(rpc("resources/list")) + assert response["result"]["resources"] == [ + {"uri": "trame://info", "name": "info", "mimeType": "application/json"}, + {"uri": "trame://logo", "name": "logo", "description": "The logo"}, + {"uri": "trame://readme", "name": "readme", "description": "How to use it."}, + {"uri": "trame://sizes", "name": "sizes"}, + ] + + async def read(uri): + response = await protocol.handle(rpc("resources/read", {"uri": uri})) + (content,) = response["result"]["contents"] + return content + + assert await read("trame://info") == { + "uri": "trame://info", + "mimeType": "application/json", + "text": '{"name": "test"}', + } + assert await read("trame://logo") == { + "uri": "trame://logo", + "mimeType": "application/octet-stream", + "blob": base64.b64encode(b"\x00\x01").decode(), + } + assert await read("trame://readme") == { + "uri": "trame://readme", + "mimeType": "text/plain", + "text": "# Readme", + } + assert (await read("trame://sizes"))["text"] == "[1, 2]" + + response = await protocol.handle(rpc("resources/read", {"uri": "trame://nope"})) + assert response["error"]["code"] == INVALID_PARAMS + + +@pytest.mark.asyncio +async def test_a_failing_resource_is_an_internal_error(mcp): + @mcp.resource("trame://broken") + def broken(): + msg = "no data" + raise RuntimeError(msg) + + response = await mcp.protocol.handle( + rpc("resources/read", {"uri": "trame://broken"}) + ) + assert response["error"] == { + "code": INTERNAL_ERROR, + "message": "RuntimeError: no data", + } + + +@pytest.mark.asyncio +async def test_resource_templates_list(protocol): + response = await protocol.handle(rpc("resources/templates/list")) + assert response["result"] == {"resourceTemplates": []} + + +@pytest.mark.asyncio +async def test_prompts(protocol, mcp): + @mcp.prompt(name="review", description="Review the session.") + async def review_session(focus: str) -> list: + return [{"role": "user", "content": {"type": "text", "text": focus}}] + + response = await protocol.handle(rpc("prompts/list")) + assert response["result"]["prompts"] == [ + { + "name": "walkthrough", + "arguments": [ + { + "name": "topic", + "required": False, + "description": "Part of the app to explain", + } + ], + }, + { + "name": "review", + "description": "Review the session.", + "arguments": [{"name": "focus", "required": True}], + }, + ] + + response = await protocol.handle( + rpc("prompts/get", {"name": "walkthrough", "arguments": {"topic": "cone"}}) + ) + assert response["result"] == { + "messages": [ + {"role": "user", "content": {"type": "text", "text": "Explain the cone"}} + ] + } + + # a list of messages is sent as is + response = await protocol.handle( + rpc("prompts/get", {"name": "review", "arguments": {"focus": "colors"}}) + ) + assert response["result"] == { + "description": "Review the session.", + "messages": [{"role": "user", "content": {"type": "text", "text": "colors"}}], + } + + response = await protocol.handle(rpc("prompts/get", {"name": "review"})) + assert response["error"]["code"] == INVALID_PARAMS + assert response["error"]["message"].startswith( + "Invalid arguments for prompt review" + ) + response = await protocol.handle(rpc("prompts/get", {"name": "nope"})) + assert response["error"] == { + "code": INVALID_PARAMS, + "message": "Unknown prompt: nope", + } + + +@pytest.mark.asyncio +async def test_disabled_blocks_tools_resources_and_prompts(protocol, mcp): + mcp.enabled = False + + response = await protocol.handle( + rpc("tools/call", {"name": "add", "arguments": {"a": 1}}) + ) + assert response["result"] == { + "content": [ + { + "type": "text", + "text": "MCP access is currently disabled by the application", + } + ], + "isError": True, + } + with pytest.raises(MCPDisabledError): + await protocol._execute(mcp.get_tool("add"), {"a": 1}) + + response = await protocol.handle(rpc("resources/read", {"uri": "trame://info"})) + assert response["error"]["code"] == INVALID_REQUEST + assert "disabled" in response["error"]["message"] + response = await protocol.handle(rpc("prompts/get", {"name": "walkthrough"})) + assert "disabled" in response["error"]["message"] + + mcp.enabled = True + assert await protocol._execute(mcp.get_tool("add"), {"a": 1}) == 2 + + +# ----------------------------------------------------------------------------- +# Running calls +# ----------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_parallel_async_tool_calls_run_one_at_a_time(mcp): + log = [] + + @mcp.tool + async def slow(tag: str) -> str: + log.append(f"start:{tag}") + await asyncio.sleep(0.01) + log.append(f"end:{tag}") + return tag + + await asyncio.gather( + call_tool(mcp, "slow", tag="a"), call_tool(mcp, "slow", tag="b") + ) + + # One at a time, each start is immediately followed by its own end + assert len(log) == 4 + for i in (0, 2): + tag = log[i].split(":")[1] + assert log[i] == f"start:{tag}" + assert log[i + 1] == f"end:{tag}" + + +@pytest.mark.asyncio +async def test_concurrent_calls_overlap(mcp): + log = [] + + @mcp.tool + async def slow(tag: str) -> str: + log.append(f"start:{tag}") + await asyncio.sleep(0.01) + log.append(f"end:{tag}") + return tag + + mcp.concurrent_calls = True + await asyncio.gather( + call_tool(mcp, "slow", tag="a"), call_tool(mcp, "slow", tag="b") + ) + assert log == ["start:a", "start:b", "end:a", "end:b"] + + +@pytest.mark.asyncio +async def test_tool_state_changes_are_flushed(mcp): + state = mcp.server.state + state.resolution = 6 + state.ready() # what starting the server does: flushes run from now on + seen = [] + + @state.change("resolution") + def on_resolution(resolution, **_): + seen.append(resolution) + + @mcp.tool + def set_resolution(resolution: int) -> int: + state.resolution = resolution # no `with state:` + return resolution + + @mcp.tool + async def set_later(resolution: int) -> int: + await asyncio.sleep(0) + state.resolution = resolution + return resolution + + await call_tool(mcp, "set_resolution", resolution=12) + assert seen == [12] + + await call_tool(mcp, "set_later", resolution=18) + assert seen == [12, 18] + + +@pytest.mark.asyncio +async def test_concurrent_tools_run_while_a_call_holds_the_lock(mcp): + release = asyncio.Event() + log = [] + + @mcp.tool + async def hold() -> str: + log.append("hold:start") + await release.wait() + log.append("hold:end") + return "held" + + @mcp.tool(read_only=True, concurrent=True) + async def peek() -> str: + log.append("peek") + return "peeked" + + @mcp.tool + def poke() -> str: + log.append("poke") + return "poked" + + holding = asyncio.ensure_future(call_tool(mcp, "hold")) + await asyncio.sleep(0) # hold takes the lock + + # The concurrent tool is not held back... + response = await asyncio.wait_for(call_tool(mcp, "peek"), 1) + assert text_of(response) == "peeked" + # ... another one waits for the lock + poking = asyncio.ensure_future(call_tool(mcp, "poke")) + await asyncio.sleep(0.01) + assert not poking.done() + + release.set() + await asyncio.gather(holding, poking) + assert log == ["hold:start", "peek", "hold:end", "poke"] + + +def test_bound_methods_carry_concurrent(mcp): + class ProtoJobs: + async def wait_for_jobs(self) -> str: + return "idle" + + def start_job(self) -> str: + return "started" + + jobs = ProtoJobs() + mcp.tool(jobs.wait_for_jobs, read_only=True, concurrent=True) + mcp.tool(jobs.start_job) + assert mcp.get_tool("wait_for_jobs").concurrent is True + assert mcp.get_tool("start_job").concurrent is False + + +@pytest.mark.asyncio +async def test_cancelled_tool_call_stops_releases_the_lock_and_flushes(mcp): + state = mcp.server.state + state.progress = 0 + state.ready() + seen, log = [], [] + + @state.change("progress") + def on_progress(progress, **_): + seen.append(progress) + + @mcp.tool + async def long_job() -> str: + state.progress = 50 + try: + await asyncio.Event().wait() # never set + finally: + log.append("cleaned up") + return "done" + + @mcp.tool + def quick() -> str: + return "quick" + + job = asyncio.ensure_future(call_tool(mcp, "long_job", rpc_id=7)) + await asyncio.sleep(0.01) + assert not job.done() + + assert mcp.protocol.cancel(8) is False # not running + assert await mcp.protocol.handle(cancel_notification(7, "user stopped")) is None + response = await asyncio.wait_for(job, 1) + assert response["id"] == 7 + assert response["error"] == { + "code": REQUEST_CANCELLED, + "message": "Request cancelled: user stopped", + } + assert log == ["cleaned up"] + assert seen == [50] # what it changed before the cancel was flushed + + # The lock is free again + response = await asyncio.wait_for(call_tool(mcp, "quick"), 1) + assert text_of(response) == "quick" + assert not mcp.protocol._in_flight + + +@pytest.mark.asyncio +async def test_cancel_reaches_only_the_requests_of_its_session(mcp): + release = asyncio.Event() + + @mcp.tool(concurrent=True) + async def wait_for_release() -> str: + await release.wait() + return "released" + + message = rpc("tools/call", {"name": "wait_for_release"}, rpc_id=1) + first = asyncio.ensure_future(mcp.protocol.handle(message, session="a")) + second = asyncio.ensure_future(mcp.protocol.handle(message, session="b")) + await asyncio.sleep(0.01) + + await mcp.protocol.handle(cancel_notification(1), session="a") + response = await asyncio.wait_for(first, 1) + assert response["error"] == { + "code": REQUEST_CANCELLED, + "message": "Request cancelled", + } + assert not second.done() + + # Unknown, finished and malformed requests are ignored + await mcp.protocol.handle(cancel_notification(1), session="a") + await mcp.protocol.handle(cancel_notification(99), session="b") + await mcp.protocol.handle(cancel_notification([1]), session="b") + await mcp.protocol.handle( + {"jsonrpc": "2.0", "method": "notifications/cancelled", "params": "x"} + ) + assert not second.done() + + release.set() + response = await asyncio.wait_for(second, 1) + assert text_of(response) == "released" + + +@pytest.mark.asyncio +async def test_a_call_cancelled_while_waiting_for_the_lock_never_runs(mcp): + release = asyncio.Event() + log = [] + + @mcp.tool + async def hold() -> str: + await release.wait() + return "held" + + @mcp.tool + def poke() -> str: + log.append("poke") + return "poked" + + holding = asyncio.ensure_future(call_tool(mcp, "hold", rpc_id=1)) + poking = asyncio.ensure_future(call_tool(mcp, "poke", rpc_id=2)) + await asyncio.sleep(0.01) + await mcp.protocol.handle(cancel_notification(2)) + assert "error" in await asyncio.wait_for(poking, 1) + + release.set() + assert "result" in await asyncio.wait_for(holding, 1) + assert log == [] + response = await asyncio.wait_for(call_tool(mcp, "poke"), 1) + assert text_of(response) == "poked" + + +@pytest.mark.asyncio +async def test_a_call_cancelled_from_outside_is_not_answered(mcp): + """The client went away (the HTTP side cancels the task): the + cancellation propagates instead of becoming a reply.""" + started = asyncio.Event() + + @mcp.tool + async def wait_forever() -> str: + started.set() + await asyncio.Event().wait() + return "done" + + @mcp.tool + def quick() -> str: + return "quick" + + handling = asyncio.ensure_future(call_tool(mcp, "wait_forever", rpc_id=3)) + await started.wait() + handling.cancel() + with pytest.raises(asyncio.CancelledError): + await handling + assert not mcp.protocol._in_flight + # The lock was released + response = await asyncio.wait_for(call_tool(mcp, "quick"), 1) + assert text_of(response) == "quick" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("method", "params"), + [ + ("tools/call", {"name": ["add"]}), + ("resources/read", {"uri": {"path": "info"}}), + ("prompts/get", {"name": ["walkthrough"]}), + ], +) +async def test_a_malformed_name_is_invalid_params(protocol, method, params): + response = await protocol.handle(rpc(method, params)) + assert response["error"]["code"] == INVALID_PARAMS diff --git a/tests/test_mcp_schema.py b/tests/test_mcp_schema.py new file mode 100644 index 0000000..20aa23e --- /dev/null +++ b/tests/test_mcp_schema.py @@ -0,0 +1,536 @@ +"""JSON Schemas from type hints, and validation of what agents send.""" + +import sys +from typing import Annotated, Literal, Optional + +import pytest +from trame_common.obj.mcp import MCPExposable + +from trame_server.mcp.schema import ( + ToolError, + ValidationError, + bind_arguments, + docstring, + from_json, + function_input_schema, + nullable, + parameter_hints, + return_hint, + type_to_schema, + unwrap, + validate, +) + + +def test_scalars(): + assert type_to_schema(str) == {"type": "string"} + assert type_to_schema(int) == {"type": "integer"} + assert type_to_schema(float) == {"type": "number"} + assert type_to_schema(bool) == {"type": "boolean"} + assert type_to_schema(type(None)) == {"type": "null"} + # each call gets its own dict + type_to_schema(str)["description"] = "changed" + assert type_to_schema(str) == {"type": "string"} + + +def test_containers(): + assert type_to_schema(list[int]) == {"type": "array", "items": {"type": "integer"}} + assert type_to_schema(set[str]) == {"type": "array", "items": {"type": "string"}} + assert type_to_schema(dict[str, float]) == { + "type": "object", + "additionalProperties": {"type": "number"}, + } + assert type_to_schema(list) == {"type": "array"} + assert type_to_schema(tuple) == {"type": "array"} + assert type_to_schema(dict) == {"type": "object"} + + +def test_optional_and_union(): + assert type_to_schema(Optional[int]) == { # noqa: UP045 - Optional is the point + "anyOf": [{"type": "integer"}, {"type": "null"}] + } + assert type_to_schema(int | None) == { + "anyOf": [{"type": "integer"}, {"type": "null"}] + } + assert type_to_schema(int | str) == { + "anyOf": [{"type": "integer"}, {"type": "string"}] + } + assert type_to_schema(int | str | None) == { + "anyOf": [{"type": "integer"}, {"type": "string"}, {"type": "null"}] + } + + +def test_literal_and_annotated(): + assert type_to_schema(Literal["a", "b"]) == {"enum": ["a", "b"]} + assert type_to_schema(Literal[1, "a", None]) == { # noqa: PYI061 + "enum": [1, "a", None] + } + assert type_to_schema(Annotated[int, "number of facets"]) == { + "type": "integer", + "description": "number of facets", + } + assert type_to_schema(Annotated[int, "facets", {"minimum": 3}]) == { + "type": "integer", + "description": "facets", + "minimum": 3, + } + + +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.10 hashes a union's members, and a dict can't be hashed", +) +def test_annotated_in_a_union(): + assert type_to_schema(Annotated[float, {"minimum": 0}] | None) == { + "anyOf": [{"type": "number", "minimum": 0}, {"type": "null"}] + } + + +def test_fixed_length_tuples(): + assert type_to_schema(tuple[float, float, float]) == { + "type": "array", + "items": {"type": "number"}, + "minItems": 3, + "maxItems": 3, + } + # JSON Schema 2020-12: positional items are prefixItems + assert type_to_schema(tuple[str, int]) == { + "type": "array", + "prefixItems": [{"type": "string"}, {"type": "integer"}], + "minItems": 2, + "maxItems": 2, + } + assert type_to_schema(tuple[int, ...]) == { + "type": "array", + "items": {"type": "integer"}, + } + + +def test_the_empty_tuple(): + assert type_to_schema(tuple[()]) == {"type": "array", "minItems": 0, "maxItems": 0} + + +def test_unknown_types_degrade_to_any(): + class Custom: + pass + + assert type_to_schema(Custom) == {} + assert type_to_schema(None) == {} + assert type_to_schema(Optional[Custom]) == {} # noqa: UP045 + + +def test_a_custom_hook_is_consulted_at_every_level(): + class Point: + pass + + def custom(annotation): + return ( + {"type": "string", "description": "a point id"} + if annotation is Point + else None + ) + + assert type_to_schema(list[Point], custom) == { + "type": "array", + "items": {"type": "string", "description": "a point id"}, + } + assert type_to_schema(dict[str, Point | None], custom)["additionalProperties"] == { + "anyOf": [{"type": "string", "description": "a point id"}, {"type": "null"}] + } + assert type_to_schema(int, custom) == {"type": "integer"} + + +def test_nullable(): + assert nullable({"type": "integer"}) == { + "anyOf": [{"type": "integer"}, {"type": "null"}] + } + for schema in ( + {}, + {"type": "null"}, + {"type": ["string", "null"]}, + {"anyOf": [{"type": "integer"}, {"type": "null"}]}, + ): + assert nullable(schema) is schema + + +def test_function_input_schema(): + def fn(resolution: int, label: str = "cone", factor: float | None = None) -> str: # noqa: ARG001 + return label + + schema = function_input_schema(fn) + assert schema["type"] == "object" + assert schema["required"] == ["resolution"] + assert schema["additionalProperties"] is False + assert schema["properties"]["resolution"] == {"type": "integer"} + assert schema["properties"]["label"] == {"type": "string", "default": "cone"} + assert schema["properties"]["factor"]["default"] is None + + +def test_function_input_schema_none_default_is_nullable(): + def fn(count: int = None): # noqa: RUF013 - implicit Optional is the point + return count + + assert function_input_schema(fn)["properties"]["count"] == { + "anyOf": [{"type": "integer"}, {"type": "null"}], + "default": None, + } + + +def test_function_input_schema_leaves_out_defaults_json_cannot_hold(): + marker = object() + + def fn(value: int = marker, items: tuple = (1, 2)): # noqa: ARG001 + return value + + properties = function_input_schema(fn)["properties"] + assert properties["value"] == {"type": "integer"} + assert properties["items"] == {"type": "array", "default": (1, 2)} + + +def test_function_input_schema_without_parameters(): + def fn(): + pass + + assert function_input_schema(fn) == { + "type": "object", + "properties": {}, + "additionalProperties": False, + } + + +def test_function_input_schema_var_keyword_allows_extras(): + def fn(a: int, *args, **options): # noqa: ARG001 + return a + + schema = function_input_schema(fn) + assert "additionalProperties" not in schema + assert list(schema["properties"]) == ["a"] + + +def test_function_input_schema_skips_self(): + class App: + def method(self, value: int): + pass + + assert list(function_input_schema(App.method)["properties"]) == ["value"] + assert list(function_input_schema(App().method)["properties"]) == ["value"] + + +def test_an_unresolvable_hint_does_not_lose_the_others(): + def fn(a: int, b: "Missing" = None) -> "AlsoMissing": # noqa: ARG001, F821 + return a + + assert parameter_hints(fn) == {"a": int} + assert return_hint(fn) is None + schema = function_input_schema(fn) + assert schema["properties"] == {"a": {"type": "integer"}, "b": {"default": None}} + + +def test_return_hint(): + def fn() -> Annotated[int, "count"]: + return 1 + + assert return_hint(fn) == Annotated[int, "count"] + + +def test_validate_scalars_and_coercion(): + assert validate(3, {"type": "integer"}) == 3 + assert validate(3.0, {"type": "integer"}) == 3 + assert isinstance(validate(3.0, {"type": "integer"}), int) + assert isinstance(validate(3, {"type": "number"}), float) + for value, schema in ( + (3.5, {"type": "integer"}), + (True, {"type": "integer"}), + (True, {"type": "number"}), + (1, {"type": "boolean"}), + ("3", {"type": "integer"}), + (None, {"type": "string"}), + (float("inf"), {"type": "integer"}), + ): + with pytest.raises(ValidationError): + validate(value, schema) + + +def test_a_validation_error_is_a_complete_tool_error(): + assert issubclass(ValidationError, ToolError) + assert issubclass(ToolError, ValueError) + with pytest.raises(ValidationError, match=r"^count: expected integer, got string$"): + validate("x", {"type": "integer"}, "count") + + +def test_validate_unknown_and_empty_schemas_accept_anything(): + assert validate({"x": [1]}, {}) == {"x": [1]} + assert validate("x", {"description": "no type", "model": "Person"}) == "x" + assert validate("x", {"type": "unknown-keyword"}) == "x" + + +def test_validate_containers_report_paths(): + schema = type_to_schema(dict[str, list[float]]) + assert validate({"a": [1, 2.5]}, schema) == {"a": [1.0, 2.5]} + with pytest.raises( + ValidationError, match=r"^value\.a\[1\]: expected number, got string$" + ): + validate({"a": [1, "x"]}, schema) + + schema = { + "type": "array", + "items": {"type": "number"}, + "minItems": 3, + "maxItems": 3, + } + with pytest.raises( + ValidationError, match=r"^color: expected at least 3 items, got 2$" + ): + validate([1, 2], schema, "color") + with pytest.raises( + ValidationError, match=r"^color: expected at most 3 items, got 4$" + ): + validate([1, 2, 3, 4], schema, "color") + + schema = type_to_schema(tuple[str, int]) + assert validate(["a", 2.0], schema) == ["a", 2] + with pytest.raises( + ValidationError, match=r"^value\[1\]: expected integer, got string$" + ): + validate(["a", "b"], schema) + + +def test_validate_objects(): + schema = { + "type": "object", + "properties": {"a": {"type": "integer"}, "b": {"type": "string"}}, + "required": ["a", "b"], + "additionalProperties": False, + } + assert validate({"a": 1.0, "b": "x"}, schema) == {"a": 1, "b": "x"} + with pytest.raises(ValidationError, match=r"^value: missing required 'a', 'b'$"): + validate({}, schema) + with pytest.raises( + ValidationError, match=r"^value: unexpected 'c' \(allowed: \['a', 'b'\]\)$" + ): + validate({"a": 1, "b": "x", "c": 3}, schema) + with pytest.raises(ValidationError, match="keys must be strings"): + validate({1: 1}, {"type": "object"}) + + extras = {"type": "object", "additionalProperties": {"type": "number"}} + assert validate({"x": 1}, extras) == {"x": 1.0} + with pytest.raises(ValidationError, match=r"^value\.x: expected number"): + validate({"x": "1"}, extras) + + +def test_validate_bounds(): + assert validate(0.5, {"type": "number", "minimum": 0, "maximum": 1}) == 0.5 + with pytest.raises(ValidationError, match=r"^value: must be at most 1, got 2\.5$"): + validate(2.5, {"type": "number", "maximum": 1}) + with pytest.raises(ValidationError, match=r"^value: must be at least 3, got 1$"): + validate(1, {"type": "integer", "minimum": 3}) + with pytest.raises( + ValidationError, match=r"^value: must be greater than 0, got 0$" + ): + validate(0, {"exclusiveMinimum": 0}) + with pytest.raises(ValidationError, match=r"^value: must be less than 1, got 1$"): + validate(1, {"exclusiveMaximum": 1}) + + +def test_validate_strings(): + schema = {"type": "string", "minLength": 2, "maxLength": 4, "pattern": "^[a-z]+$"} + assert validate("abc", schema) == "abc" + with pytest.raises(ValidationError, match="at least 2 characters"): + validate("a", schema) + with pytest.raises(ValidationError, match="at most 4 characters"): + validate("abcde", schema) + with pytest.raises(ValidationError, match="must match pattern"): + validate("ABC", schema) + + +def test_enum_and_const_quote_values_as_json(): + with pytest.raises(ValidationError) as error: + validate("cubes", {"enum": ["surface", "wireframe"]}, "arguments.mode") + assert str(error.value) == ( + 'arguments.mode: must be one of "surface", "wireframe", got "cubes"' + ) + with pytest.raises(ValidationError) as error: + validate(3, {"enum": [1, None, True]}) + assert str(error.value) == "value: must be one of 1, null, true, got 3" + with pytest.raises(ValidationError) as error: + validate("b", {"const": "a"}) + assert str(error.value) == 'value: must be "a"' + + assert validate("a", {"const": "a"}) == "a" + assert validate(1, {"enum": [1, 2]}) == 1 + # True must not match 1, as in JSON + with pytest.raises(ValidationError): + validate(True, {"enum": [1, 2]}) + with pytest.raises(ValidationError): + validate(1, {"const": True}) + + +def test_validate_unions(): + schema = type_to_schema(int | str) + assert validate("a", schema) == "a" + with pytest.raises(ValidationError) as error: + validate([1], schema) + assert str(error.value) == ( + "value: does not match any allowed form " + "(expected integer, got array; expected string, got array)" + ) + assert validate(None, type_to_schema(Optional[int])) is None # noqa: UP045 + assert validate(2, {"oneOf": [{"type": "string"}, {"type": "integer"}]}) == 2 + + assert validate("on", {"type": ["boolean", "string"]}) == "on" + with pytest.raises( + ValidationError, match=r"^value: expected boolean or string, got integer$" + ): + validate(1, {"type": ["boolean", "string"]}) + + +def test_bind_arguments_checks_schema_and_signature(): + def fn(a: int, b: str = "x"): # noqa: ARG001 + return a + + schema = function_input_schema(fn) + assert bind_arguments(fn, {"a": 1.0}, schema) == {"a": 1} + assert bind_arguments(lambda: None, None, {}) == {} + with pytest.raises( + ValidationError, match=r"^arguments: expected object, got array$" + ): + bind_arguments(fn, [1], schema) + with pytest.raises(ValidationError, match="unexpected 'c'"): + bind_arguments(fn, {"a": 1, "c": 2}, schema) + + # a permissive hand-written schema still cannot smuggle unknown arguments + with pytest.raises(ValidationError, match="unexpected keyword argument"): + bind_arguments(fn, {"a": 1, "c": 2}, {"type": "object"}) + with pytest.raises(ValidationError, match=r"^args: missing a required argument"): + bind_arguments(fn, {}, {"type": "object"}, label="args") + + +def test_bind_arguments_takes_leading_positional_arguments(): + def build(server, size: int): + return server, size + + schema = function_input_schema(build, skip=("server",)) + assert bind_arguments(build, {"size": 2}, schema, positional=("srv",)) == { + "size": 2 + } + with pytest.raises(ValidationError, match="missing a required argument: 'size'"): + bind_arguments(build, {}, {"type": "object"}, positional=("srv",)) + + +def test_from_json_restores_declared_tuples(): + Vector = tuple[float, float, float] + assert from_json([1.0, 2.0, 3.0], Vector) == (1.0, 2.0, 3.0) + assert from_json([[0, 1], [2, 3]], list[tuple[int, int]]) == [(0, 1), (2, 3)] + assert from_json([1, 2, 3], tuple[int, ...]) == (1, 2, 3) + assert from_json([1, 2], tuple) == (1, 2) + assert from_json({"a": [1, 2]}, dict[str, tuple[int, int]]) == {"a": (1, 2)} + assert from_json([1.0, 2.0, 3.0], Vector | None) == (1.0, 2.0, 3.0) + assert from_json(None, Vector | None) is None + assert from_json([1.0, 2.0, 3.0], Annotated[Vector, "a point"]) == (1.0, 2.0, 3.0) + assert from_json([1.0, 2.0, 3.0], Annotated[Vector, "a point"] | None) == ( + 1.0, + 2.0, + 3.0, + ) + assert from_json([1.0, 2.0, 3.0], Optional[Annotated[Vector, "p"]]) == ( # noqa: UP045 + 1.0, + 2.0, + 3.0, + ) + assert from_json({"a": [1, 2]}, dict[str, tuple[int, int]] | None) == {"a": (1, 2)} + # anything else is left alone + assert from_json([1, 2], list[int]) == [1, 2] + assert from_json([1, 2], list) == [1, 2] + assert from_json("x", Vector) == "x" + assert from_json("x", Vector | str) == "x" + assert from_json([1, 2], None) == [1, 2] + + +def test_tool_arguments_get_their_declared_tuples(): + def place(position: tuple[float, float, float], tags: list[str] | None = None): + return position, tags + + schema = function_input_schema(place) + arguments = bind_arguments(place, {"position": [1, 2, 3], "tags": ["a"]}, schema) + assert arguments == {"position": (1.0, 2.0, 3.0), "tags": ["a"]} + assert isinstance(arguments["position"], tuple) + + +def test_from_json_restores_declared_sets(): + assert from_json(["a", "b", "a"], set[str]) == {"a", "b"} + assert from_json([1, 2], frozenset[int]) == frozenset({1, 2}) + assert from_json([1, 2], set) == {1, 2} + assert from_json({"k": ["a"]}, dict[str, set[str]] | None) == {"k": {"a"}} + + +class Thing(MCPExposable): + pass + + +def test_from_json_resolves_ids_where_objects_are_declared(): + things = {"0x1": "one", "0x2": "two"} + seen = [] + + def resolve(value, cls, path): + seen.append((value, cls, path)) + return things[value] + + assert from_json("0x1", Thing, resolve) == "one" + assert from_json(["0x1", "0x2"], list[Thing], resolve, "args.items") == [ + "one", + "two", + ] + assert seen[-1] == ("0x2", Thing, "args.items[1]") + assert from_json({"a": "0x2"}, dict[str, Thing] | None, resolve, "v") == { + "a": "two" + } + assert seen[-1] == ("0x2", Thing, "v.a") + assert from_json("0x1", Annotated[Thing, "The thing"] | None, resolve) == "one" + # Without resolve, ids stay as given; a second pass resolves them + shaped = from_json(["0x1"], set[Thing]) + assert shaped == {"0x1"} + assert from_json(shaped, set[Thing], resolve) == {"one"} + + +# ----------------------------------------------------------------------------- +# Docstrings +# ----------------------------------------------------------------------------- + + +def test_wrapped_lines_are_joined_paragraphs_kept(): + assert unwrap("Orient the camera so\nthis axis points up.") == ( + "Orient the camera so this axis points up." + ) + assert unwrap("One that\nwraps.\n\nTwo that\nwraps.") == ( + "One that wraps.\n\nTwo that wraps." + ) + assert unwrap(None) is None + + +def test_lists_fields_and_indented_blocks_keep_their_breaks(): + assert unwrap("Modes:\n- fast, which\n wraps\n- slow\n1. one\n2) two") == ( + "Modes:\n- fast, which wraps\n- slow\n1. one\n2) two" + ) + assert unwrap("Load.\n\n:param path: the path\n to it\n:returns: it") == ( + "Load.\n\n:param path: the path to it\n:returns: it" + ) + assert unwrap("Args:\n x: the x\n y: the y\nThen prose\ngoes on.") == ( + "Args:\n x: the x\n y: the y\nThen prose goes on." + ) + assert unwrap("Run:\n```\na = 1\nb = 2\n```\nafter the\nfence.") == ( + "Run:\n```\na = 1\nb = 2\n```\nafter the fence." + ) + + +def test_docstring(): + def documented(): + """Reset the camera so + everything shows. + + Keeps the view up. + """ + + assert docstring(documented) == ( + "Reset the camera so everything shows.\n\nKeeps the view up." + ) + assert docstring(lambda: None) is None