From d603615d6a4cdefcac3dc217f21c424056b5b467 Mon Sep 17 00:00:00 2001 From: worldmozara Date: Tue, 8 Sep 2026 12:06:34 +0800 Subject: [PATCH] tools: register MCP server tools in the catalog (mcp source, purge) --- src/plyngent/tools/catalog.py | 13 +++- src/plyngent/tools/mcp.py | 108 +++++++++++++++++++++++++++++ tests/test_tools/test_mcp_tools.py | 93 +++++++++++++++++++++++++ 3 files changed, 211 insertions(+), 3 deletions(-) create mode 100644 src/plyngent/tools/mcp.py create mode 100644 tests/test_tools/test_mcp_tools.py diff --git a/src/plyngent/tools/catalog.py b/src/plyngent/tools/catalog.py index 725f00c..01d37dd 100644 --- a/src/plyngent/tools/catalog.py +++ b/src/plyngent/tools/catalog.py @@ -12,7 +12,7 @@ from plyngent.agent.tools import ToolDefinition, ToolTag if TYPE_CHECKING: from collections.abc import Container, Generator -type ToolSourceKind = Literal["builtin", "plugin"] +type ToolSourceKind = Literal["builtin", "plugin", "mcp"] type ToolSurface = Literal["local", "public"] @@ -29,8 +29,8 @@ class ToolSource: def __str__(self) -> str: if self.kind == "builtin": return "builtin" - plugin = self.plugin_id or "?" - return f"plugin:{plugin}" + id_part = self.plugin_id or "?" + return f"{self.kind}:{id_part}" @dataclass(frozen=True, slots=True) @@ -152,6 +152,13 @@ class ToolCatalog: def clear(self) -> None: self._by_name.clear() + def purge_kind(self, kind: ToolSourceKind) -> int: + """Remove every entry registered from *kind*; return the count removed.""" + doomed = [name for name, entry in self._by_name.items() if entry.source.kind == kind] + for name in doomed: + del self._by_name[name] + return len(doomed) + def snapshot(self) -> dict[str, RegisteredTool]: return dict(self._by_name) diff --git a/src/plyngent/tools/mcp.py b/src/plyngent/tools/mcp.py new file mode 100644 index 0000000..6ab5f94 --- /dev/null +++ b/src/plyngent/tools/mcp.py @@ -0,0 +1,108 @@ +"""Register MCP server tools into the tool catalog (namespaced, LOCAL tags). + +Mirrors :mod:`plyngent.tools.plugins`: the catalog only *registers* what the +model may see; hosts still select into a :class:`ToolRegistry`. Connections +must already be started (see :class:`~plyngent.runtime.mcp_client.McpManager`); +this module is sync-only so registry rebuilds never need an event loop. + +Tool names are namespaced ``mcp____`` so two servers (or a +server and a builtin) can never collide on sanitized names; the catalog still +refuses true clashes. +""" + +from __future__ import annotations + +import re +from typing import TYPE_CHECKING + +from plyngent.agent.tools import ToolDefinition, ToolTag +from plyngent.tools.catalog import ToolSource, get_catalog + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + + from plyngent.config.models import McpConfig + from plyngent.runtime.mcp_client import McpManager, McpTool + +_MCP_PREFIX = "mcp__" +_NAME_SEGMENT_MAX = 64 +_SANITIZE_PATTERN = re.compile(r"[^A-Za-z0-9_-]") + + +def _sanitize_segment(part: str) -> str: + """Reduce a server/tool name to a model-safe, bounded name segment.""" + return _SANITIZE_PATTERN.sub("_", part)[:_NAME_SEGMENT_MAX] or "srv" + + +def mcp_tool_name(server: str, tool: str) -> str: + """Namespaced catalog name for one MCP tool (collision-free per server).""" + return f"{_MCP_PREFIX}{_sanitize_segment(server)}__{_sanitize_segment(tool)}" + + +def _make_handler( + manager: McpManager, + server: str, + tool: str, +) -> Callable[..., Awaitable[str]]: + async def handler(**arguments: object) -> str: + return await manager.call(server, tool, dict(arguments)) + + return handler + + +def register_mcp_server_tools( + server: str, + tools: list[McpTool], + *, + manager: McpManager, + read_only: bool = False, +) -> list[str]: + """Register one server's advertised tools; returns the registered names. + + Tags: LOCAL (plus READ_ONLY when the server config opts in) — never YOLO or + TRUSTABLE, so host-side danger classification always sees every call. + """ + catalog = get_catalog() + source = ToolSource(kind="mcp", plugin_id=server) + names: list[str] = [] + for entry in tools: + if not entry.name: + continue + tags = ToolTag.LOCAL + if read_only: + tags |= ToolTag.READ_ONLY + description = entry.description or f"MCP tool {entry.name!r} from server {server!r}" + definition = ToolDefinition( + name=mcp_tool_name(server, entry.name), + description=f"[mcp:{server}] {description}", + parameters=dict(entry.input_schema), + handler=_make_handler(manager, server, entry.name), + tags=tags, + ) + catalog.register(definition, source=source) + names.append(definition.name) + return names + + +def purge_mcp_tools() -> int: + """Remove previously registered MCP tools (call before re-registering).""" + return get_catalog().purge_kind("mcp") + + +def register_mcp_tools(manager: McpManager, config: McpConfig) -> list[str]: + """Register every connected server's tools; replaces prior MCP entries.""" + _ = purge_mcp_tools() + names: list[str] = [] + for connection in manager.connections(): + if not connection.connected or not connection.tools: + continue + server_cfg = config.servers.get(connection.name) + names.extend( + register_mcp_server_tools( + connection.name, + connection.tools, + manager=manager, + read_only=server_cfg.read_only if server_cfg is not None else False, + ) + ) + return names diff --git a/tests/test_tools/test_mcp_tools.py b/tests/test_tools/test_mcp_tools.py new file mode 100644 index 0000000..05d3a9e --- /dev/null +++ b/tests/test_tools/test_mcp_tools.py @@ -0,0 +1,93 @@ +"""MCP tool registration: namespacing, tags, catalog source, purge, handlers.""" + +from __future__ import annotations + +import pytest + +from plyngent.agent.tools import ToolTag +from plyngent.runtime.mcp_client import McpTool +from plyngent.tools.catalog import catalog_scope +from plyngent.tools.mcp import mcp_tool_name, purge_mcp_tools, register_mcp_server_tools + + +class _StubManager: + """Duck-typed McpManager stand-in (records calls, returns canned text).""" + + def __init__(self) -> None: + self.calls: list[tuple[str, str, dict[str, object]]] = [] + + async def call(self, server: str, tool: str, arguments: dict[str, object]) -> str: + self.calls.append((server, tool, arguments)) + return f"{server}:{tool}:{arguments.get('text', '')}" + + +def _tools() -> list[McpTool]: + return [ + McpTool( + name="echo", + description="Echo text.", + input_schema={"type": "object", "properties": {"text": {"type": "string"}}}, + ), + McpTool(name="", description="skipped: empty name", input_schema={}), + ] + + +def test_mcp_tool_name_sanitizes_and_namespaces() -> None: + assert mcp_tool_name("my server", "a/b c") == "mcp__my_server__a_b_c" + assert mcp_tool_name("x", "y").startswith("mcp__") + + +def test_register_names_tags_source_and_purge() -> None: + manager = _StubManager() + with catalog_scope(empty=True): + names = register_mcp_server_tools("docs", _tools(), manager=manager) # type: ignore[arg-type] + assert names == ["mcp__docs__echo"] + + from plyngent.tools.catalog import get_catalog + + entry = get_catalog().get("mcp__docs__echo") + assert entry is not None + assert entry.source.kind == "mcp" + assert entry.source.plugin_id == "docs" + assert str(entry.source) == "mcp:docs" + definition = entry.definition + assert definition.tags == ToolTag.LOCAL + assert definition.description.startswith("[mcp:docs]") + assert definition.parameters["type"] == "object" + assert "text" in definition.parameters["properties"] + + assert purge_mcp_tools() == 1 + assert get_catalog().get("mcp__docs__echo") is None + assert purge_mcp_tools() == 0 + + +def test_register_read_only_flag() -> None: + manager = _StubManager() + with catalog_scope(empty=True): + register_mcp_server_tools("docs", _tools()[:1], manager=manager, read_only=True) # type: ignore[arg-type] + from plyngent.tools.catalog import get_catalog + + entry = get_catalog().get("mcp__docs__echo") + assert entry is not None + assert entry.definition.tags == ToolTag.LOCAL | ToolTag.READ_ONLY + + +def test_duplicate_registration_collides() -> None: + manager = _StubManager() + with catalog_scope(empty=True): + register_mcp_server_tools("docs", _tools()[:1], manager=manager) # type: ignore[arg-type] + with pytest.raises(ValueError, match="collision"): + register_mcp_server_tools("docs", _tools()[:1], manager=manager) # type: ignore[arg-type] + + +async def test_handler_proxies_to_manager_call() -> None: + manager = _StubManager() + with catalog_scope(empty=True): + register_mcp_server_tools("docs", _tools()[:1], manager=manager) # type: ignore[arg-type] + from plyngent.tools.catalog import get_catalog + + entry = get_catalog().get("mcp__docs__echo") + assert entry is not None + result = await entry.definition.handler(text="hi") # type: ignore[union-attr] + assert result == "docs:echo:hi" + assert manager.calls == [("docs", "echo", {"text": "hi"})]