monitor_app/mcp/rucio.py

0001 """JLab and BNL Rucio catalogs on the authenticated SWF MCP service.
0002
0003 The maintained rucio-eic MCP implementation is loaded twice under isolated
0004 module names because it reads one catalog configuration at import time.  The
0005 plain tool functions are then registered on swf-monitor's FastMCP instance
0006 with catalog prefixes.  Credentials, token caches, and the BNL X509 proxy stay
0007 on swf-testbed; remote MCP clients see only read-only catalog operations.
0008 """
0009
0010 from __future__ import annotations
0011
0012 import asyncio
0013 import importlib.util
0014 import inspect
0015 import os
0016 from pathlib import Path
0017 from types import ModuleType
0018
0019 from django.conf import settings
0020
0021 from monitor_app.mcp import mcp
0022
0023
0024 _CATALOGS = {
0025     'jlab_rucio': {
0026         'description': (
0027             'JLab Rucio science-data catalog (normally scope epic, with '
0028             'path-like /RECO and /SIMU dataset names). Read-only.'),
0029         'env': {
0030             'RUCIO_URL': settings.RUCIO_JLAB_URL,
0031             'RUCIO_AUTH_TYPE': 'userpass',
0032             'RUCIO_ACCOUNT': settings.RUCIO_JLAB_ACCOUNT,
0033             'RUCIO_USERNAME': settings.RUCIO_JLAB_USERNAME,
0034             'RUCIO_PASSWORD': settings.RUCIO_JLAB_PASSWORD,
0035             'RUCIO_VO': '',
0036             'RUCIO_CA_BUNDLE': 'false',
0037             'TOKEN_FILE_PATH': settings.RUCIO_JLAB_TOKEN_FILE,
0038             'X509_USER_PROXY': '',
0039         },
0040     },
0041     'bnl_rucio': {
0042         'description': (
0043             'BNL Rucio PanDA production catalog (normally scope group.EIC). '
0044             'Use for output and log registration, rules, locks, replicas, '
0045             'and RSE diagnostics. Read-only.'),
0046         'env': {
0047             'RUCIO_URL': settings.RUCIO_BNL_URL,
0048             'RUCIO_AUTH_TYPE': 'x509',
0049             'RUCIO_ACCOUNT': settings.RUCIO_BNL_ACCOUNT,
0050             'RUCIO_USERNAME': '',
0051             'RUCIO_PASSWORD': '',
0052             'RUCIO_VO': settings.RUCIO_BNL_VO,
0053             'RUCIO_CA_BUNDLE': settings.RUCIO_BNL_CA_BUNDLE,
0054             'TOKEN_FILE_PATH': settings.RUCIO_BNL_TOKEN_FILE,
0055             'X509_USER_PROXY': settings.RUCIO_BNL_X509_PROXY,
0056         },
0057     },
0058 }
0059
0060
0061 def _load_catalog(prefix: str, values: dict[str, str]) -> ModuleType:
0062     """Load one independently configured copy of the Rucio MCP module."""
0063     path = Path(settings.RUCIO_MCP_MODULE_PATH)
0064     if not path.is_file():
0065         raise RuntimeError(f'Rucio MCP module not found: {path}')
0066
0067     previous = {key: os.environ.get(key) for key in values}
0068     try:
0069         os.environ.update({key: str(value) for key, value in values.items()})
0070         spec = importlib.util.spec_from_file_location(
0071             f'_swf_monitor_{prefix}', path)
0072         if spec is None or spec.loader is None:
0073             raise RuntimeError(f'Cannot load Rucio MCP module: {path}')
0074         module = importlib.util.module_from_spec(spec)
0075         spec.loader.exec_module(module)
0076         return module
0077     finally:
0078         for key, value in previous.items():
0079             if value is None:
0080                 os.environ.pop(key, None)
0081             else:
0082                 os.environ[key] = value
0083
0084
0085 def _register_catalog(prefix: str, config: dict) -> list[str]:
0086     module = _load_catalog(prefix, config['env'])
0087     registered = []
0088     # Mirror the complete upstream service. This is deliberately discovery-
0089     # driven rather than a local tool allowlist; new Rucio MCP capabilities
0090     # appear here automatically after that maintained service is updated.
0091     upstream_tools = module.mcp._tool_manager.list_tools()
0092     for upstream_tool in upstream_tools:
0093         function_name = upstream_tool.name
0094         source = getattr(module, function_name)
0095         tool_name = f'{prefix}_{function_name}'
0096
0097         async def wrapper(_source=source, **kwargs):
0098             return await asyncio.to_thread(_source, **kwargs)
0099
0100         wrapper.__name__ = tool_name
0101         wrapper.__doc__ = inspect.getdoc(source) or ''
0102         wrapper.__signature__ = inspect.signature(source)
0103         source_description = upstream_tool.description or (
0104             wrapper.__doc__.splitlines()[0] if wrapper.__doc__ else '')
0105         description = f"{config['description']} {source_description}".strip()
0106         mcp.tool(
0107             name=tool_name,
0108             description=description,
0109         )(wrapper)
0110         RUCIO_TOOL_DISCOVERY.append({
0111             'name': tool_name,
0112             'description': description,
0113             'parameters': list(inspect.signature(source).parameters),
0114         })
0115         registered.append(tool_name)
0116     return registered
0117
0118
0119 RUCIO_TOOL_DISCOVERY = []
0120 RUCIO_TOOL_NAMES = []
0121 for _prefix, _config in _CATALOGS.items():
0122     RUCIO_TOOL_NAMES.extend(_register_catalog(_prefix, _config))
0123
0124
0125 def get_rucio_tool_discovery() -> list[dict]:
0126     """Discovery records for swf_list_available_tools."""
0127     return [dict(record) for record in RUCIO_TOOL_DISCOVERY]