@register_per_user_function_group(config_type=DataRobotMCPClientConfig)
async def datarobot_mcp_client_function_group(
config: DataRobotMCPClientConfig, _builder: Builder
) -> MCPFunctionGroup:
"""
Connect to an MCP server and expose tools as a function group.
Args:
config: The configuration for the MCP client
_builder: The builder
Returns:
The function group
"""
from nat.plugins.mcp.client.client_base import MCPSSEClient # noqa: PLC0415
from nat.plugins.mcp.client.client_impl import MCPFunctionGroup # noqa: PLC0415
from nat.plugins.mcp.client.client_impl import (
mcp_apply_tool_alias_and_description, # noqa: PLC0415
)
from nat.plugins.mcp.client.client_impl import mcp_session_tool_function # noqa: PLC0415
# Resolve auth provider if specified
auth_provider = None
if config.server.auth_provider:
auth_provider = await _builder.get_auth_provider(config.server.auth_provider)
# Build the appropriate client
if config.server.transport == "sse":
client = MCPSSEClient(
str(config.server.url),
tool_call_timeout=config.tool_call_timeout,
auth_flow_timeout=config.auth_flow_timeout,
reconnect_enabled=config.reconnect_enabled,
reconnect_max_attempts=config.reconnect_max_attempts,
reconnect_initial_backoff=config.reconnect_initial_backoff,
reconnect_max_backoff=config.reconnect_max_backoff,
)
elif config.server.transport == "streamable-http":
# Use default_user_id for the base client
# For interactive OAuth2: from config. For service accounts: defaults to server URL
base_user_id = (
getattr(auth_provider.config, "default_user_id", str(config.server.url))
if auth_provider
else None
)
client = DataRobotMCPStreamableHTTPClient(
str(config.server.url),
auth_provider=auth_provider,
user_id=base_user_id,
tool_call_timeout=config.tool_call_timeout,
auth_flow_timeout=config.auth_flow_timeout,
reconnect_enabled=config.reconnect_enabled,
reconnect_max_attempts=config.reconnect_max_attempts,
reconnect_initial_backoff=config.reconnect_initial_backoff,
reconnect_max_backoff=config.reconnect_max_backoff,
)
else:
raise ValueError(f"Unsupported transport: {config.server.transport}")
logger.info("Configured to use MCP server at %s", client.server_name)
# Create the MCP function group
group = MCPFunctionGroup(config=config)
# Store shared components for session client creation
group._shared_auth_provider = auth_provider
group._client_config = config
# Set auth provider config defaults
# For interactive OAuth2: use config values
# For service accounts: default_user_id = server URL,
# allow_default_user_id_for_tool_calls = True
if auth_provider:
group._default_user_id = getattr(
auth_provider.config, "default_user_id", str(config.server.url)
)
group._allow_default_user_id_for_tool_calls = getattr(
auth_provider.config, "allow_default_user_id_for_tool_calls", True
)
else:
group._default_user_id = None
group._allow_default_user_id_for_tool_calls = True
yielded = False
try:
async with client:
# Expose the live MCP client on the function group instance so other components
# (e.g., HTTP endpoints) can reuse the already-established session instead of creating a
# new client per request.
group.mcp_client = client
group.mcp_client_server_name = client.server_name
group.mcp_client_transport = client.transport
all_tools = await client.get_tools()
tool_overrides = mcp_apply_tool_alias_and_description(all_tools, config.tool_overrides)
# Add each tool as a function to the group
for tool_name, tool in all_tools.items():
# Get override if it exists
override = tool_overrides.get(tool_name)
# Use override values or defaults
function_name = override.alias if override and override.alias else tool_name
description = (
override.description if override and override.description else tool.description
)
# Create the tool function according to configuration.
# Patch the input schema to store enum values as strings so
# NAT's downstream model_validate(kwargs) never trips on
# cross-class enum identity (BUZZOK-30556).
tool_fn = _make_input_schema_enum_safe(mcp_session_tool_function(tool, group))
# Normalize optional typing for linter/type-checker compatibility
single_fn = tool_fn.single_fn
if single_fn is None:
# Should not happen because FunctionInfo always sets a single_fn
logger.warning("Skipping tool %s because single_fn is None", function_name)
continue
input_schema = tool_fn.input_schema
# Convert NoneType sentinel to None for FunctionGroup.add_function signature
if input_schema is type(None): # noqa: E721
input_schema = None
# Add to group
logger.debug("Adding tool %s to group", function_name)
group.add_function(
name=function_name,
description=description,
fn=single_fn,
input_schema=input_schema,
converters=tool_fn.converters,
)
yielded = True
yield group
except Exception as e:
if hasattr(e, "exceptions"):
primary_exception = extract_primary_exception(list(e.exceptions))
else:
primary_exception = e
logger.warning("Error in MCP client function group: %s", primary_exception)
group.mcp_client = None
group.mcp_client_server_name = getattr(client, "server_name", str(config.server.url))
group.mcp_client_transport = getattr(client, "transport", config.server.transport)
if not yielded:
yield group
else:
# Cleanup (e.g. __aexit__) failed after we already yielded; re-raise so
# the caller sees it and we do not yield a second time.
raise