ryan
0b8fe7b2df
Load an allowlisted subset of bookstack-mcp tools (search/read plus create/update books, chapters, pages and comments; no deletes) over stdio and add them to the admin agent alongside Homebox. Extract the shared allowlist + is_admin wrapping into utils/mcp_tools.py. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
112 lines
3.7 KiB
Python
112 lines
3.7 KiB
Python
"""Shared helpers for loading admin-only agent tools from stdio MCP servers.
|
|
|
|
Each integration (Homebox, BookStack, ...) picks an allowlist of server tools;
|
|
they are renamed ``<prefix>_<name>``, get trimmed descriptions, and re-check
|
|
``is_admin`` in the run config on every call as a second line of defense
|
|
(the tools are only compiled into the admin agent in the first place).
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
|
|
from langchain_core.runnables import RunnableConfig
|
|
from langchain_core.tools import BaseTool, StructuredTool
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
DESCRIPTION_MAX_CHARS = 300
|
|
|
|
|
|
def content_to_text(content) -> str:
|
|
"""Flatten MCP content blocks into plain text for the LLM and persistence."""
|
|
if isinstance(content, str):
|
|
return content
|
|
if isinstance(content, (list, tuple)):
|
|
parts = []
|
|
for block in content:
|
|
if isinstance(block, dict):
|
|
parts.append(block.get("text") or f"[{block.get('type', 'content')}]")
|
|
else:
|
|
parts.append(str(block))
|
|
return "\n".join(parts)
|
|
return str(content)
|
|
|
|
|
|
def short_description(description: str) -> str:
|
|
description = " ".join((description or "").split())
|
|
if len(description) > DESCRIPTION_MAX_CHARS:
|
|
description = description[: DESCRIPTION_MAX_CHARS - 3].rstrip() + "..."
|
|
return description
|
|
|
|
|
|
def admin_only_message(label: str) -> str:
|
|
return f"{label} access is restricted to admin users."
|
|
|
|
|
|
def wrap_admin_only(inner: BaseTool, prefix: str, label: str) -> StructuredTool:
|
|
"""Rename an MCP tool to ``<prefix>_<name>`` and gate it on ``is_admin``."""
|
|
|
|
async def call(config: RunnableConfig, **kwargs) -> str:
|
|
if not (config or {}).get("configurable", {}).get("is_admin"):
|
|
return admin_only_message(label)
|
|
try:
|
|
result = await inner.coroutine(**kwargs)
|
|
except Exception as e:
|
|
return f"{label} error: {e}"
|
|
content = result[0] if isinstance(result, tuple) else result
|
|
return content_to_text(content)
|
|
|
|
return StructuredTool(
|
|
name=f"{prefix}_{inner.name}",
|
|
description=short_description(inner.description),
|
|
args_schema=inner.args_schema,
|
|
coroutine=call,
|
|
metadata=inner.metadata,
|
|
)
|
|
|
|
|
|
def select_tools(
|
|
mcp_tools: list[BaseTool], allowlist: tuple[str, ...], prefix: str, label: str
|
|
) -> list[StructuredTool]:
|
|
"""Keep allowlisted tools (in allowlist order) and wrap them admin-only."""
|
|
by_name = {t.name: t for t in mcp_tools}
|
|
missing = [name for name in allowlist if name not in by_name]
|
|
if missing:
|
|
logger.warning(f"{label} MCP: server is missing tools {missing}")
|
|
return [
|
|
wrap_admin_only(by_name[name], prefix, label)
|
|
for name in allowlist
|
|
if name in by_name
|
|
]
|
|
|
|
|
|
async def load_stdio_tools(
|
|
server: str,
|
|
command: str,
|
|
env: dict[str, str],
|
|
allowlist: tuple[str, ...],
|
|
prefix: str,
|
|
label: str,
|
|
) -> list[StructuredTool]:
|
|
"""Start a stdio MCP server and return its wrapped tools, or [] on failure."""
|
|
try:
|
|
from langchain_mcp_adapters.client import MultiServerMCPClient
|
|
|
|
client = MultiServerMCPClient(
|
|
{
|
|
server: {
|
|
"transport": "stdio",
|
|
"command": command,
|
|
"args": [],
|
|
"env": {**env, "PATH": os.environ.get("PATH", "")},
|
|
}
|
|
}
|
|
)
|
|
tools = select_tools(await client.get_tools(), allowlist, prefix, label)
|
|
except Exception:
|
|
logger.exception(f"{label} MCP failed to load; continuing without it")
|
|
return []
|
|
|
|
logger.info(f"{label} MCP: loaded {len(tools)} tools")
|
|
return tools
|