mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-07 23:47:54 +08:00
tools: register MCP server tools in the catalog (mcp source, purge)
This commit is contained in:
@@ -12,7 +12,7 @@ from plyngent.agent.tools import ToolDefinition, ToolTag
|
|||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from collections.abc import Container, Generator
|
from collections.abc import Container, Generator
|
||||||
|
|
||||||
type ToolSourceKind = Literal["builtin", "plugin"]
|
type ToolSourceKind = Literal["builtin", "plugin", "mcp"]
|
||||||
type ToolSurface = Literal["local", "public"]
|
type ToolSurface = Literal["local", "public"]
|
||||||
|
|
||||||
|
|
||||||
@@ -29,8 +29,8 @@ class ToolSource:
|
|||||||
def __str__(self) -> str:
|
def __str__(self) -> str:
|
||||||
if self.kind == "builtin":
|
if self.kind == "builtin":
|
||||||
return "builtin"
|
return "builtin"
|
||||||
plugin = self.plugin_id or "?"
|
id_part = self.plugin_id or "?"
|
||||||
return f"plugin:{plugin}"
|
return f"{self.kind}:{id_part}"
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True, slots=True)
|
@dataclass(frozen=True, slots=True)
|
||||||
@@ -152,6 +152,13 @@ class ToolCatalog:
|
|||||||
def clear(self) -> None:
|
def clear(self) -> None:
|
||||||
self._by_name.clear()
|
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]:
|
def snapshot(self) -> dict[str, RegisteredTool]:
|
||||||
return dict(self._by_name)
|
return dict(self._by_name)
|
||||||
|
|
||||||
|
|||||||
@@ -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__<server>__<tool>`` 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
|
||||||
@@ -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"})]
|
||||||
Reference in New Issue
Block a user