nvidia-nat-mcp 1.3.0a20250910__py3-none-any.whl → 1.3.0a20250922__py3-none-any.whl

nat/plugins/mcp/client_base.py

@@ -20,12 +20,9 @@ from abc import ABC
 from abc import abstractmethod
 from contextlib import AsyncExitStack
 from contextlib import asynccontextmanager
-from enum import Enum
-from typing import Any
+from typing import AsyncGenerator
 
-from pydantic import BaseModel
-from pydantic import Field
-from pydantic import create_model
+import httpx
 
 from mcp import ClientSession
 from mcp.client.sse import sse_client
@@ -33,104 +30,120 @@ from mcp.client.stdio import StdioServerParameters
 from mcp.client.stdio import stdio_client
 from mcp.client.streamable_http import streamablehttp_client
 from mcp.types import TextContent
+from nat.authentication.interfaces import AuthProviderBase
+from nat.data_models.authentication import AuthReason
+from nat.data_models.authentication import AuthRequest
 from nat.plugins.mcp.exception_handler import mcp_exception_handler
 from nat.plugins.mcp.exceptions import MCPToolNotFoundError
+from nat.plugins.mcp.utils import model_from_mcp_schema
 from nat.utils.type_utils import override
 
 logger = logging.getLogger(__name__)
 
 
-def model_from_mcp_schema(name: str, mcp_input_schema: dict) -> type[BaseModel]:
+class AuthAdapter(httpx.Auth):
     """
-    Create a pydantic model from the input schema of the MCP tool
+    httpx.Auth adapter for authentication providers.
+    Converts AuthProviderBase to httpx.Auth interface for dynamic token management.
     """
-    _type_map = {
-        "string": str,
-        "number": float,
-        "integer": int,
-        "boolean": bool,
-        "array": list,
-        "null": None,
-        "object": dict,
-    }
-
-    properties = mcp_input_schema.get("properties", {})
-    required_fields = set(mcp_input_schema.get("required", []))
-    schema_dict = {}
-
-    def _generate_valid_classname(class_name: str):
-        return class_name.replace('_', ' ').replace('-', ' ').title().replace(' ', '')
-
-    def _generate_field(field_name: str, field_properties: dict[str, Any]) -> tuple:
-        json_type = field_properties.get("type", "string")
-        enum_vals = field_properties.get("enum")
-
-        if enum_vals:
-            enum_name = f"{field_name.capitalize()}Enum"
-            field_type = Enum(enum_name, {item: item for item in enum_vals})
-
-        elif json_type == "object" and "properties" in field_properties:
-            field_type = model_from_mcp_schema(name=field_name, mcp_input_schema=field_properties)
-        elif json_type == "array" and "items" in field_properties:
-            item_properties = field_properties.get("items", {})
-            if item_properties.get("type") == "object":
-                item_type = model_from_mcp_schema(name=field_name, mcp_input_schema=item_properties)
+
+    def __init__(self, auth_provider: AuthProviderBase, auth_for_tool_calls_only: bool = False):
+        self.auth_provider = auth_provider
+        self.auth_for_tool_calls_only = auth_for_tool_calls_only
+
+    async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]:
+        """Add authentication headers to the request using NAT auth provider."""
+        # Check if we should only auth tool calls, Is this needed?
+        if self.auth_for_tool_calls_only and not self._is_tool_call_request(request):
+            # Skip auth for non-tool calls
+            yield request
+            return
+
+        try:
+            # Get fresh auth headers from the NAT auth provider
+            auth_headers = await self._get_auth_headers(reason=AuthReason.NORMAL)
+            request.headers.update(auth_headers)
+        except Exception as e:
+            logger.info("Failed to get auth headers: %s", e)
+            # Continue without auth headers if auth fails
+
+        response = yield request
+
+        # Handle 401 responses by retrying with fresh auth
+        if response.status_code == 401:
+            try:
+                # Get fresh auth headers with 401 context
+                auth_headers = await self._get_auth_headers(reason=AuthReason.RETRY_AFTER_401, response=response)
+                request.headers.update(auth_headers)
+                yield request  # Retry the request
+            except Exception as e:
+                logger.info("Failed to refresh auth after 401: %s", e)
+        return
+
+    def _is_tool_call_request(self, request: httpx.Request) -> bool:
+        """Check if this is a tool call request based on the request body."""
+        try:
+            # Check if the request body contains a tool call
+            if request.content:
+                import json
+                body = json.loads(request.content.decode('utf-8'))
+                # Check if it's a JSON-RPC request with method "tools/call"
+                if (isinstance(body, dict) and body.get("method") == "tools/call"):
+                    return True
+        except (json.JSONDecodeError, UnicodeDecodeError, AttributeError):
+            # If we can't parse the body, assume it's not a tool call
+            pass
+        return False
+
+    async def _get_auth_headers(self, reason: AuthReason, response: httpx.Response | None = None) -> dict[str, str]:
+        """Get authentication headers from the NAT auth provider."""
+        # Build auth request
+        www_authenticate = response.headers.get("WWW-Authenticate", None) if response else None
+        auth_request = AuthRequest(
+            reason=reason,
+            www_authenticate=www_authenticate,
+        )
+        try:
+            # Mutating the config is not thread-safe, so we need to lock here
+            # Is mutating the config the only way to pass the auth request to the auth provider? This needs
+            # to be re-visited.
+            self.auth_provider.config.auth_request = auth_request
+            auth_result = await self.auth_provider.authenticate()
+            # Check if we have BearerTokenCred
+            from nat.data_models.authentication import BearerTokenCred
+            if auth_result.credentials and isinstance(auth_result.credentials[0], BearerTokenCred):
+                token = auth_result.credentials[0].token.get_secret_value()
+                return {"Authorization": f"Bearer {token}"}
             else:
-                item_type = _type_map.get(item_properties.get("type", "string"), Any)
-            field_type = list[item_type]
-        elif isinstance(json_type, list):
-            field_type = None
-            for t in json_type:
-                mapped = _type_map.get(t, Any)
-                field_type = mapped if field_type is None else field_type | mapped
-
-            return field_type, Field(
-                default=field_properties.get("default", None if "null" in json_type else ...),
-                description=field_properties.get("description", "")
-            )
-        else:
-            field_type = _type_map.get(json_type, Any)
-
-        # Determine the default value based on whether the field is required
-        if field_name in required_fields:
-            # Field is required - use explicit default if provided, otherwise make it required
-            default_value = field_properties.get("default", ...)
-        else:
-            # Field is optional - use explicit default if provided, otherwise None
-            default_value = field_properties.get("default", None)
-            # Make the type optional if no default was provided
-            if "default" not in field_properties:
-                field_type = field_type | None
-
-        nullable = field_properties.get("nullable", False)
-        description = field_properties.get("description", "")
-
-        field_type = field_type | None if nullable else field_type
-
-        return field_type, Field(default=default_value, description=description)
-
-    for field_name, field_props in properties.items():
-        schema_dict[field_name] = _generate_field(field_name=field_name, field_properties=field_props)
-    return create_model(f"{_generate_valid_classname(name)}InputSchema", **schema_dict)
+                logger.warning("Auth provider did not return BearerTokenCred")
+                return {}
+        except Exception as e:
+            logger.warning("Failed to get auth token: %s", e)
+            return {}
 
 
 class MCPBaseClient(ABC):
     """
-    Base client for creating a session and connecting to an MCP server
+    Base client for creating a MCP transport session and connecting to an MCP server
 
     Args:
         transport (str): The type of client to use ('sse', 'stdio', or 'streamable-http')
+        auth_provider (AuthProviderBase | None): Optional authentication provider for Bearer token injection
     """
 
-    def __init__(self, transport: str = 'streamable-http'):
+    def __init__(self, transport: str = 'streamable-http', auth_provider: AuthProviderBase | None = None):
         self._tools = None
         self._transport = transport.lower()
         if self._transport not in ['sse', 'stdio', 'streamable-http']:
             raise ValueError("transport must be either 'sse', 'stdio' or 'streamable-http'")
 
         self._exit_stack: AsyncExitStack | None = None
+        self._session: ClientSession | None = None  # Main session
+        self._connection_established = False
+        self._initial_connection = False
 
-        self._session: ClientSession | None = None
+        # Convert auth provider to AuthAdapter
+        self._httpx_auth = AuthAdapter(auth_provider) if auth_provider else None
 
     @property
     def transport(self) -> str:
@@ -142,15 +155,19 @@ class MCPBaseClient(ABC):
 
         self._exit_stack = AsyncExitStack()
 
+        # Establish connection with httpx.Auth
         self._session = await self._exit_stack.enter_async_context(self.connect_to_server())
 
+        self._initial_connection = True
+        self._connection_established = True
+
         return self
 
     async def __aexit__(self, exc_type, exc_value, traceback):
-
         if not self._exit_stack:
             raise RuntimeError("MCPBaseClient not initialized. Use async with to initialize.")
 
+        # Close session
         await self._exit_stack.aclose()
         self._session = None
         self._exit_stack = None
@@ -168,11 +185,12 @@ class MCPBaseClient(ABC):
         """
         Establish a session with an MCP server within an async context
         """
-        pass
+        yield
 
     async def get_tools(self):
         """
         Retrieve a dictionary of all tools served by the MCP server.
+        Uses unauthenticated session for discovery.
         """
 
         if not self._session:
@@ -185,7 +203,8 @@ class MCPBaseClient(ABC):
                 MCPToolClient(session=self._session,
                               tool_name=tool.name,
                               tool_description=tool.description,
-                              tool_input_schema=tool.inputSchema)
+                              tool_input_schema=tool.inputSchema,
+                              parent_client=self)
             for tool in response.tools
         }
 
@@ -257,7 +276,9 @@ class MCPSSEClient(MCPBaseClient):
 
 class MCPStdioClient(MCPBaseClient):
     """
-    Client for creating a session and connecting to an MCP server using stdio
+    Client for creating a session and connecting to an MCP server using stdio.
+    This is a local transport that spawns the MCP server process and communicates
+    with it over stdin/stdout.
 
     Args:
       command (str): The command to run
@@ -307,11 +328,11 @@ class MCPStreamableHTTPClient(MCPBaseClient):
 
     Args:
       url (str): The url of the MCP server
+      auth_provider (AuthProviderBase | None): Optional authentication provider for Bearer token injection
     """
 
-    def __init__(self, url: str):
-        super().__init__("streamable-http")
-
+    def __init__(self, url: str, auth_provider: AuthProviderBase | None = None):
+        super().__init__("streamable-http", auth_provider=auth_provider)
         self._url = url
 
     @property
@@ -323,11 +344,13 @@ class MCPStreamableHTTPClient(MCPBaseClient):
         return f"streamable-http:{self._url}"
 
     @asynccontextmanager
+    @override
     async def connect_to_server(self):
         """
         Establish a session with an MCP server via streamable-http within an async context
         """
-        async with streamablehttp_client(url=self._url) as (read, write, get_session_id):
+        # Use httpx.Auth for authentication
+        async with streamablehttp_client(url=self._url, auth=self._httpx_auth) as (read, write, _):
             async with ClientSession(read, write) as session:
                 await session.initialize()
                 yield session
@@ -335,24 +358,28 @@ class MCPStreamableHTTPClient(MCPBaseClient):
 
 class MCPToolClient:
     """
-    Client wrapper used to call an MCP tool.
+    Client wrapper used to call an MCP tool. This assumes that the MCP transport session
+    has already been setup.
 
     Args:
-        connect_fn (callable): Function that returns an async context manager for connecting to the server
+        session (ClientSession): The MCP client session
         tool_name (str): The name of the tool to wrap
         tool_description (str): The description of the tool provided by the MCP server.
         tool_input_schema (dict): The input schema for the tool.
+        parent_client (MCPBaseClient): The parent MCP client for auth management.
     """
 
     def __init__(self,
                  session: ClientSession,
                  tool_name: str,
                  tool_description: str | None,
-                 tool_input_schema: dict | None = None):
+                 tool_input_schema: dict | None = None,
+                 parent_client: "MCPBaseClient | None" = None):
         self._session = session
         self._tool_name = tool_name
         self._tool_description = tool_description
         self._input_schema = (model_from_mcp_schema(self._tool_name, tool_input_schema) if tool_input_schema else None)
+        self._parent_client = parent_client
 
     @property
     def name(self):
@@ -388,6 +415,9 @@ class MCPToolClient:
         Args:
             tool_args (dict[str, Any]): A dictionary of key value pairs to serve as inputs for the MCP tool.
         """
+        if self._session is None:
+            raise RuntimeError("No session available for tool call")
+        logger.info("Calling tool %s with arguments %s", self._tool_name, tool_args)
         result = await self._session.call_tool(self._tool_name, tool_args)
 
         output = []

nat/plugins/mcp/client_impl.py

@@ -22,11 +22,11 @@ from pydantic import HttpUrl
 from pydantic import model_validator
 
 from nat.builder.builder import Builder
-from nat.builder.function_info import FunctionInfo
-from nat.cli.register_workflow import register_function
-from nat.data_models.function import FunctionBaseConfig
-from nat.experimental.decorators.experimental_warning_decorator import experimental
-from nat.plugins.mcp.client_base import MCPBaseClient
+from nat.builder.function import FunctionGroup
+from nat.cli.register_workflow import register_function_group
+from nat.data_models.component_ref import AuthenticationRef
+from nat.data_models.function import FunctionGroupBaseConfig
+from nat.plugins.mcp.tool import mcp_tool_function
 
 logger = logging.getLogger(__name__)
 
@@ -54,6 +54,9 @@ class MCPServerConfig(BaseModel):
     args: list[str] | None = Field(default=None, description="Arguments for the stdio command")
     env: dict[str, str] | None = Field(default=None, description="Environment variables for the stdio process")
 
+    # Authentication configuration
+    auth_provider: AuthenticationRef | None = Field(default=None, description="Reference to authentication provider")
+
     @model_validator(mode="after")
     def validate_model(self):
         """Validate that stdio and SSE/Streamable HTTP properties are mutually exclusive."""
@@ -62,168 +65,118 @@ class MCPServerConfig(BaseModel):
                 raise ValueError("url should not be set when using stdio transport")
             if not self.command:
                 raise ValueError("command is required when using stdio transport")
-        elif self.transport in ("sse", "streamable-http"):
+            # Auth is not supported for stdio transport
+            if self.auth_provider is not None:
+                raise ValueError("Authentication is not supported for stdio transport")
+        elif self.transport == "sse":
+            if self.command is not None or self.args is not None or self.env is not None:
+                raise ValueError("command, args, and env should not be set when using sse transport")
+            if not self.url:
+                raise ValueError("url is required when using sse transport")
+            # Auth is not supported for SSE transport
+            if self.auth_provider is not None:
+                raise ValueError("Authentication is not supported for SSE transport.")
+        elif self.transport == "streamable-http":
             if self.command is not None or self.args is not None or self.env is not None:
-                raise ValueError("command, args, and env should not be set when using sse or streamable-http transport")
+                raise ValueError("command, args, and env should not be set when using streamable-http transport")
             if not self.url:
-                raise ValueError("url is required when using sse or streamable-http transport")
+                raise ValueError("url is required when using streamable-http transport")
+
         return self
 
 
-class MCPClientConfig(FunctionBaseConfig, name="mcp_client"):
+class MCPClientConfig(FunctionGroupBaseConfig, name="mcp_client"):
     """
     Configuration for connecting to an MCP server as a client and exposing selected tools.
     """
     server: MCPServerConfig = Field(..., description="Server connection details (transport, url/command, etc.)")
-    tool_filter: dict[str, MCPToolOverrideConfig] | list[str] | None = Field(
+    tool_overrides: dict[str, MCPToolOverrideConfig] | None = Field(
         default=None,
-        description="""Filter or map tools to expose from the server (list or dict).
-        Can be:
-        - A list of tool names to expose: ['tool1', 'tool2']
-        - A dict mapping tool names to override configs:
-          {'tool1': {'alias': 'new_name', 'description': 'New desc'}}
-          {'tool2': {'description': 'Override description only'}}  # alias defaults to 'tool2'
+        description="""Optional tool name overrides and description changes.
+        Example:
+          tool_overrides:
+            calculator_add:
+              alias: "add_numbers"
+              description: "Add two numbers together"
+            calculator_multiply:
+              description: "Multiply two numbers"  # alias defaults to original name
         """)
 
 
-class MCPSingleToolConfig(FunctionBaseConfig, name="mcp_single_tool"):
-    """
-    Configuration for wrapping a single tool from an MCP server as a NeMo Agent toolkit function.
-    """
-    client: MCPBaseClient = Field(..., description="MCP client to use for the tool")
-    tool_name: str = Field(..., description="Name of the tool to use")
-    tool_description: str | None = Field(default=None, description="Description of the tool")
-
-    model_config = {"arbitrary_types_allowed": True}
-
-
-def _get_server_name_safe(client: MCPBaseClient) -> str:
-    # Avoid leaking env secrets from stdio client in logs.
-    if client.transport == "stdio":
-        safe_server = f"stdio: {client.command}"
-    else:
-        safe_server = f"{client.transport}: {client.url}"
-
-    return safe_server
-
-
-@register_function(config_type=MCPSingleToolConfig)
-async def mcp_single_tool(config: MCPSingleToolConfig, builder: Builder):
+@register_function_group(config_type=MCPClientConfig)
+async def mcp_client_function_group(config: MCPClientConfig, _builder: Builder):
     """
-    Wrap a single tool from an MCP server as a NeMo Agent toolkit function.
-    """
-    tool = await config.client.get_tool(config.tool_name)
-    if config.tool_description:
-        tool.set_description(description=config.tool_description)
-    input_schema = tool.input_schema
-
-    logger.info("Configured to use tool: %s from MCP server at %s", tool.name, _get_server_name_safe(config.client))
-
-    def _convert_from_str(input_str: str) -> BaseModel:
-        return input_schema.model_validate_json(input_str)
-
-    @experimental(feature_name="mcp_client")
-    async def _response_fn(tool_input: BaseModel | None = None, **kwargs) -> str:
-        try:
-            if tool_input:
-                return await tool.acall(tool_input.model_dump())
-            _ = input_schema.model_validate(kwargs)
-            return await tool.acall(kwargs)
-        except Exception as e:
-            return str(e)
-
-    fn = FunctionInfo.create(single_fn=_response_fn,
-                             description=tool.description,
-                             input_schema=input_schema,
-                             converters=[_convert_from_str])
-    yield fn
-
-
-@register_function(MCPClientConfig)
-async def mcp_client_function_handler(config: MCPClientConfig, builder: Builder):
-    """
-    Connect to an MCP server, discover tools, and register them as functions in the workflow.
-
-    Note:
-    - Uses builder's exit stack to manage client lifecycle
-    - Applies tool filters if provided
+    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_base import MCPSSEClient
     from nat.plugins.mcp.client_base import MCPStdioClient
     from nat.plugins.mcp.client_base import MCPStreamableHTTPClient
 
-    # Build the appropriate client
-    client_cls = {
-        "stdio": lambda: MCPStdioClient(config.server.command, config.server.args, config.server.env),
-        "sse": lambda: MCPSSEClient(str(config.server.url)),
-        "streamable-http": lambda: MCPStreamableHTTPClient(str(config.server.url)),
-    }.get(config.server.transport)
+    # Resolve auth provider if specified
+    auth_provider = None
+    if config.server.auth_provider:
+        auth_provider = await _builder.get_auth_provider(config.server.auth_provider)
 
-    if not client_cls:
+    # Build the appropriate client
+    if config.server.transport == "stdio":
+        if not config.server.command:
+            raise ValueError("command is required for stdio transport")
+        client = MCPStdioClient(config.server.command, config.server.args, config.server.env)
+    elif config.server.transport == "sse":
+        client = MCPSSEClient(str(config.server.url))
+    elif config.server.transport == "streamable-http":
+        client = MCPStreamableHTTPClient(str(config.server.url), auth_provider=auth_provider)
+    else:
         raise ValueError(f"Unsupported transport: {config.server.transport}")
 
-    client = client_cls()
-    logger.info("Configured to use MCP server at %s", _get_server_name_safe(client))
+    logger.info("Configured to use MCP server at %s", client.server_name)
+
+    # Create the function group
+    group = FunctionGroup(config=config)
 
-    # client aenter connects to the server and stores the client in the exit stack
-    # so it's cleaned up when the workflow is done
     async with client:
         all_tools = await client.get_tools()
-        tool_configs = mcp_filter_tools(all_tools, config.tool_filter)
+        tool_overrides = mcp_apply_tool_alias_and_description(all_tools, config.tool_overrides)
 
-        for tool_name, tool_cfg in tool_configs.items():
-            await builder.add_function(
-                tool_cfg["function_name"],
-                MCPSingleToolConfig(
-                    client=client,
-                    tool_name=tool_name,
-                    tool_description=tool_cfg["description"],
-                ))
+        # 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)
 
-        @experimental(feature_name="mcp_client")
-        async def idle_fn(text: str) -> str:
-            # This function is a placeholder and will be removed when function groups are used
-            return f"MCP client connected: {text}"
+            # 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
 
-        yield FunctionInfo.create(single_fn=idle_fn, description="MCP client")
+            # Create the tool function
+            tool_fn = mcp_tool_function(tool)
 
+            # Add to group
+            logger.info("Adding tool %s to group", function_name)
+            group.add_function(name=function_name,
+                               description=description,
+                               fn=tool_fn.single_fn,
+                               input_schema=tool_fn.input_schema,
+                               converters=tool_fn.converters)
 
-def mcp_filter_tools(all_tools: dict, tool_filter) -> dict[str, dict]:
-    """
-    Apply tool filtering and optional aliasing/description overrides.
+        yield group
 
+
+def mcp_apply_tool_alias_and_description(
+        all_tools: dict, tool_overrides: dict[str, MCPToolOverrideConfig] | None) -> dict[str, MCPToolOverrideConfig]:
+    """
+    Filter tool overrides to only include tools that exist in the MCP server.
+    Args:
+        all_tools: The tools from the MCP server
+        tool_overrides: The tool overrides to apply
     Returns:
-        Dict[str, dict] where each value has:
-            - function_name
-            - description
+        Dictionary of valid tool overrides
     """
-    if tool_filter is None:
-        return {name: {"function_name": name, "description": tool.description} for name, tool in all_tools.items()}
-
-    if isinstance(tool_filter, list):
-        return {
-            name: {
-                "function_name": name, "description": all_tools[name].description
-            }
-            for name in tool_filter if name in all_tools
-        }
-
-    if isinstance(tool_filter, dict):
-        result = {}
-        for name, override in tool_filter.items():
-            tool = all_tools.get(name)
-            if not tool:
-                logger.warning("Tool '%s' specified in tool_filter not found in MCP server", name)
-                continue
-
-            if isinstance(override, MCPToolOverrideConfig):
-                result[name] = {
-                    "function_name": override.alias or name, "description": override.description or tool.description
-                }
-            else:
-                logger.warning("Unsupported override type for '%s': %s", name, type(override))
-                result[name] = {"function_name": name, "description": tool.description}
-        return result
-
-    # Fallback for unsupported tool_filter types
-    raise ValueError(f"Unsupported tool_filter type: {type(tool_filter)}")
+    if not tool_overrides:
+        return {}
+
+    return {name: override for name, override in tool_overrides.items() if name in all_tools}