feat(mcp): serve agents (MCP) from the trame server

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",
 ]

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:
         """

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",
+]

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