nvidia-nat-mcp 1.3.0rc1__py3-none-any.whl → 1.4.0a20251008__py3-none-any.whl

nat/plugins/mcp/client_base.py

@@ -16,7 +16,6 @@
 from __future__ import annotations
 
 import asyncio
-import json
 import logging
 from abc import ABC
 from abc import abstractmethod
@@ -55,10 +54,13 @@ class AuthAdapter(httpx.Auth):
     Converts AuthProviderBase to httpx.Auth interface for dynamic token management.
     """
 
-    def __init__(self, auth_provider: AuthProviderBase):
+    def __init__(self, auth_provider: AuthProviderBase, user_id: str | None = None):
         self.auth_provider = auth_provider
+        self.user_id = user_id  # Session-specific user ID for cache isolation
         # each adapter instance has its own lock to avoid unnecessary delays for multiple clients
         self._lock = anyio.Lock()
+        # Track whether we're currently in an interactive authentication flow
+        self.is_authenticating = False
 
     async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]:
         """Add authentication headers to the request using NAT auth provider."""
@@ -85,48 +87,30 @@ class AuthAdapter(httpx.Auth):
                     # 4. The auth headers are revoked
                     # 5. Auth config on the MCP server has changed
                     # In this case we attempt to re-run discovery and authentication
+
+                    # Signal that we're entering interactive auth flow
+                    self.is_authenticating = True
+                    logger.debug("Starting authentication flow due to 401 response")
+
                     auth_headers = await self._get_auth_headers(request=request, 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)
+                    raise
+                finally:
+                    # Signal that auth flow is complete
+                    self.is_authenticating = False
+                    logger.debug("Authentication flow completed")
         return
 
-    def _get_session_id_from_tool_call_request(self, request: httpx.Request) -> tuple[str | None, bool]:
-        """Check if this is a tool call request based on the request body.
-        Return the session id if it exists and a boolean indicating if it is a tool call request
-        """
-        try:
-            # Check if the request body contains a tool call
-            if request.content:
-                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"):
-                    session_id = body.get("params").get("_meta").get("session_id")
-                    return session_id, True
-        except (json.JSONDecodeError, UnicodeDecodeError, AttributeError):
-            # If we can't parse the body, assume it's not a tool call
-            pass
-        return None, False
-
     async def _get_auth_headers(self,
                                 request: httpx.Request | None = None,
                                 response: httpx.Response | None = None) -> dict[str, str]:
         """Get authentication headers from the NAT auth provider."""
         try:
-            session_id = None
-            is_tool_call = False
-            if request:
-                session_id, is_tool_call = self._get_session_id_from_tool_call_request(request)
-
-            if is_tool_call:
-                # Tool call requests should use the session id
-                user_id = session_id
-            else:
-                # Non-tool call requests should use the session id if it exists and fallback to default user id
-                user_id = session_id or self.auth_provider.config.default_user_id
-
-            auth_result = await self.auth_provider.authenticate(user_id=user_id, response=response)
+            # Use the user_id passed to this AuthAdapter instance
+            auth_result = await self.auth_provider.authenticate(user_id=self.user_id, response=response)
 
             # Check if we have BearerTokenCred
             from nat.data_models.authentication import BearerTokenCred
@@ -148,12 +132,20 @@ class MCPBaseClient(ABC):
     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
+        tool_call_timeout (timedelta): Timeout for tool calls when authentication is not required
+        auth_flow_timeout (timedelta): Extended timeout for tool calls that may require interactive authentication
+        reconnect_enabled (bool): Whether to automatically reconnect on connection failures
+        reconnect_max_attempts (int): Maximum number of reconnection attempts
+        reconnect_initial_backoff (float): Initial backoff delay in seconds for reconnection attempts
+        reconnect_max_backoff (float): Maximum backoff delay in seconds for reconnection attempts
     """
 
     def __init__(self,
                  transport: str = 'streamable-http',
                  auth_provider: AuthProviderBase | None = None,
-                 tool_call_timeout: timedelta = timedelta(seconds=5),
+                 user_id: str | None = None,
+                 tool_call_timeout: timedelta = timedelta(seconds=60),
+                 auth_flow_timeout: timedelta = timedelta(seconds=300),
                  reconnect_enabled: bool = True,
                  reconnect_max_attempts: int = 2,
                  reconnect_initial_backoff: float = 0.5,
@@ -170,9 +162,12 @@ class MCPBaseClient(ABC):
 
         # Convert auth provider to AuthAdapter
         self._auth_provider = auth_provider
-        self._httpx_auth = AuthAdapter(auth_provider) if auth_provider else None
+        # Use provided user_id or fall back to auth provider's default_user_id
+        effective_user_id = user_id or (auth_provider.config.default_user_id if auth_provider else None)
+        self._httpx_auth = AuthAdapter(auth_provider, effective_user_id) if auth_provider else None
 
         self._tool_call_timeout = tool_call_timeout
+        self._auth_flow_timeout = auth_flow_timeout
 
         # Reconnect configuration
         self._reconnect_enabled = reconnect_enabled
@@ -267,12 +262,25 @@ class MCPBaseClient(ABC):
     async def _with_reconnect(self, coro):
         """
         Execute an awaited operation, reconnecting once on errors.
+        Does not reconnect if the error occurs during an active authentication flow.
         """
         try:
             return await coro()
         except Exception as e:
+            # Check if error happened during active authentication flow
+            if self._httpx_auth and self._httpx_auth.is_authenticating:
+                # Provide specific error message for authentication timeouts
+                if isinstance(e, TimeoutError):
+                    logger.error("Timeout during user authentication flow - user may have abandoned authentication")
+                    raise RuntimeError(
+                        "Authentication timed out. User did not complete authentication in browser within "
+                        f"{self._auth_flow_timeout.total_seconds()} seconds.") from e
+                else:
+                    logger.error("Error during authentication flow: %s", e)
+                    raise
+
+            # Normal error - attempt reconnect if enabled
             if self._reconnect_enabled:
-                logger.warning("MCP Client operation failed. Attempting reconnect: %s", e)
                 try:
                     await self._reconnect()
                 except Exception as reconnect_err:
@@ -281,6 +289,47 @@ class MCPBaseClient(ABC):
                 return await coro()
             raise
 
+    async def _has_cached_auth_token(self) -> bool:
+        """
+        Check if we have a cached, non-expired authentication token.
+
+        Returns:
+            bool: True if we have a valid cached token, False if authentication may be needed
+        """
+        if not self._auth_provider:
+            return True  # No auth needed
+
+        try:
+            # Check if OAuth2 provider has tokens cached
+            if hasattr(self._auth_provider, '_auth_code_provider'):
+                provider = self._auth_provider._auth_code_provider
+                if provider and hasattr(provider, '_authenticated_tokens'):
+                    # Check if we have at least one non-expired token
+                    for auth_result in provider._authenticated_tokens.values():
+                        if not auth_result.is_expired():
+                            return True
+
+            return False
+        except Exception:
+            # If we can't check, assume we need auth to be safe
+            return False
+
+    async def _get_tool_call_timeout(self) -> timedelta:
+        """
+        Determine the appropriate timeout for a tool call based on authentication state.
+
+        Returns:
+            timedelta: auth_flow_timeout if authentication may be needed, tool_call_timeout otherwise
+        """
+        if self._auth_provider:
+            has_token = await self._has_cached_auth_token()
+            timeout = self._tool_call_timeout if has_token else self._auth_flow_timeout
+            if not has_token:
+                logger.debug("Using extended timeout (%s) for potential interactive authentication", timeout)
+            return timeout
+        else:
+            return self._tool_call_timeout
+
     @mcp_exception_handler
     async def get_tools(self) -> dict[str, MCPToolClient]:
         """
@@ -313,8 +362,7 @@ class MCPBaseClient(ABC):
                               tool_name=tool.name,
                               tool_description=tool.description,
                               tool_input_schema=tool.inputSchema,
-                              parent_client=self,
-                              tool_call_timeout=self._tool_call_timeout)
+                              parent_client=self)
             for tool in response.tools
         }
 
@@ -348,35 +396,13 @@ class MCPBaseClient(ABC):
         if self._auth_provider and hasattr(self._auth_provider, "_set_custom_auth_callback"):
             self._auth_provider._set_custom_auth_callback(auth_callback)
 
-    @mcp_exception_handler
-    async def call_tool_with_meta(self, tool_name: str, args: dict, session_id: str):
-        from mcp.types import CallToolRequest
-        from mcp.types import CallToolRequestParams
-        from mcp.types import CallToolResult
-        from mcp.types import ClientRequest
-
-        if not self._session:
-            raise RuntimeError("MCPBaseClient not initialized. Use async with to initialize.")
-
-        async def _call_tool_with_meta():
-            params = CallToolRequestParams(name=tool_name, arguments=args, **{"_meta": {"session_id": session_id}})
-            req = ClientRequest(CallToolRequest(params=params))
-            # We will increase the timeout to 5 minutes if the tool call timeout is less than 5 min and
-            # auth is enabled.
-            if self._auth_provider and self._tool_call_timeout.total_seconds() < 300:
-                timeout = timedelta(seconds=300)
-            else:
-                timeout = self._tool_call_timeout
-            return await self._session.send_request(req, CallToolResult, request_read_timeout_seconds=timeout)
-
-        return await self._with_reconnect(_call_tool_with_meta)
-
     @mcp_exception_handler
     async def call_tool(self, tool_name: str, tool_args: dict | None):
 
         async def _call_tool():
             session = self._session
-            return await session.call_tool(tool_name, tool_args, read_timeout_seconds=self._tool_call_timeout)
+            timeout = await self._get_tool_call_timeout()
+            return await session.call_tool(tool_name, tool_args, read_timeout_seconds=timeout)
 
         return await self._with_reconnect(_call_tool)
 
@@ -391,13 +417,15 @@ class MCPSSEClient(MCPBaseClient):
 
     def __init__(self,
                  url: str,
-                 tool_call_timeout: timedelta = timedelta(seconds=5),
+                 tool_call_timeout: timedelta = timedelta(seconds=60),
+                 auth_flow_timeout: timedelta = timedelta(seconds=300),
                  reconnect_enabled: bool = True,
                  reconnect_max_attempts: int = 2,
                  reconnect_initial_backoff: float = 0.5,
                  reconnect_max_backoff: float = 50.0):
         super().__init__("sse",
                          tool_call_timeout=tool_call_timeout,
+                         auth_flow_timeout=auth_flow_timeout,
                          reconnect_enabled=reconnect_enabled,
                          reconnect_max_attempts=reconnect_max_attempts,
                          reconnect_initial_backoff=reconnect_initial_backoff,
@@ -440,13 +468,15 @@ class MCPStdioClient(MCPBaseClient):
                  command: str,
                  args: list[str] | None = None,
                  env: dict[str, str] | None = None,
-                 tool_call_timeout: timedelta = timedelta(seconds=5),
+                 tool_call_timeout: timedelta = timedelta(seconds=60),
+                 auth_flow_timeout: timedelta = timedelta(seconds=300),
                  reconnect_enabled: bool = True,
                  reconnect_max_attempts: int = 2,
                  reconnect_initial_backoff: float = 0.5,
                  reconnect_max_backoff: float = 50.0):
         super().__init__("stdio",
                          tool_call_timeout=tool_call_timeout,
+                         auth_flow_timeout=auth_flow_timeout,
                          reconnect_enabled=reconnect_enabled,
                          reconnect_max_attempts=reconnect_max_attempts,
                          reconnect_initial_backoff=reconnect_initial_backoff,
@@ -497,14 +527,18 @@ class MCPStreamableHTTPClient(MCPBaseClient):
     def __init__(self,
                  url: str,
                  auth_provider: AuthProviderBase | None = None,
-                 tool_call_timeout: timedelta = timedelta(seconds=5),
+                 user_id: str | None = None,
+                 tool_call_timeout: timedelta = timedelta(seconds=60),
+                 auth_flow_timeout: timedelta = timedelta(seconds=300),
                  reconnect_enabled: bool = True,
                  reconnect_max_attempts: int = 2,
                  reconnect_initial_backoff: float = 0.5,
                  reconnect_max_backoff: float = 50.0):
         super().__init__("streamable-http",
                          auth_provider=auth_provider,
+                         user_id=user_id,
                          tool_call_timeout=tool_call_timeout,
+                         auth_flow_timeout=auth_flow_timeout,
                          reconnect_enabled=reconnect_enabled,
                          reconnect_max_attempts=reconnect_max_attempts,
                          reconnect_initial_backoff=reconnect_initial_backoff,
@@ -550,14 +584,12 @@ class MCPToolClient:
                  parent_client: MCPBaseClient,
                  tool_name: str,
                  tool_description: str | None,
-                 tool_input_schema: dict | None = None,
-                 tool_call_timeout: timedelta = timedelta(seconds=5)):
+                 tool_input_schema: dict | 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
-        self._tool_call_timeout = tool_call_timeout
 
         if self._parent_client is None:
             raise RuntimeError("MCPToolClient initialized without a parent client.")
@@ -589,35 +621,10 @@ class MCPToolClient:
         """
         self._tool_description = description
 
-    def _get_session_id(self) -> str | None:
-        """
-        Get the session id from the context.
-        """
-        from nat.builder.context import Context as _Ctx
-
-        # get auth callback (for example: WebSocketAuthenticationFlowHandler). this is lazily set in the client
-        # on first tool call
-        auth_callback = _Ctx.get().user_auth_callback
-        if auth_callback and self._parent_client:
-            # set custom auth callback
-            self._parent_client.set_user_auth_callback(auth_callback)
-
-        # get session id from context, authentication is done per-websocket session for tool calls
-        session_id = None
-        cookies = getattr(_Ctx.get().metadata, "cookies", None)
-        if cookies:
-            session_id = cookies.get("nat-session")
-
-        if not session_id:
-            # use default user id if allowed
-            if self._parent_client.auth_provider and \
-                self._parent_client.auth_provider.config.allow_default_user_id_for_tool_calls:
-                session_id = self._parent_client.auth_provider.config.default_user_id
-        return session_id
-
     async def acall(self, tool_args: dict) -> str:
         """
         Call the MCP tool with the provided arguments.
+        Session context is now handled at the client level, eliminating the need for metadata injection.
 
         Args:
             tool_args (dict[str, Any]): A dictionary of key value pairs to serve as inputs for the MCP tool.
@@ -625,25 +632,10 @@ class MCPToolClient:
         if self._session is None:
             raise RuntimeError("No session available for tool call")
 
-        # Extract context information
-        try:
-            session_id = self._get_session_id()
-        except Exception:
-            session_id = None
-
         try:
-            # if auth is enabled and session id is not available return user is not authorized to call the tool
-            if self._parent_client.auth_provider and not session_id:
-                result_str = "User is not authorized to call the tool"
-                mcp_error: MCPError = convert_to_mcp_error(RuntimeError(result_str), self._parent_client.server_name)
-                raise mcp_error
-
-            if session_id:
-                logger.info("Calling tool %s with arguments %s for a user session", self._tool_name, tool_args)
-                result = await self._parent_client.call_tool_with_meta(self._tool_name, tool_args, session_id)
-            else:
-                logger.info("Calling tool %s with arguments %s", self._tool_name, tool_args)
-                result = await self._session.call_tool(self._tool_name, tool_args)
+            # Simple tool call - session context is already in the client instance
+            logger.info("Calling tool %s with arguments %s", self._tool_name, tool_args)
+            result = await self._parent_client.call_tool(self._tool_name, tool_args)
 
             output = []
             for res in result.content:

nat/plugins/mcp/client_config.py

@@ -0,0 +1,131 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from datetime import timedelta
+from typing import Literal
+
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic import HttpUrl
+from pydantic import model_validator
+
+from nat.data_models.component_ref import AuthenticationRef
+from nat.data_models.function import FunctionGroupBaseConfig
+
+
+class MCPToolOverrideConfig(BaseModel):
+    """
+    Configuration for overriding tool properties when exposing from MCP server.
+    """
+    alias: str | None = Field(default=None, description="Override the tool name (function name in the workflow)")
+    description: str | None = Field(default=None, description="Override the tool description")
+
+
+class MCPServerConfig(BaseModel):
+    """
+    Server connection details for MCP client.
+    Supports stdio, sse, and streamable-http transports.
+    streamable-http is the recommended default for HTTP-based connections.
+    """
+    transport: Literal["stdio", "sse", "streamable-http"] = Field(
+        ..., description="Transport type to connect to the MCP server (stdio, sse, or streamable-http)")
+    url: HttpUrl | None = Field(default=None,
+                                description="URL of the MCP server (for sse or streamable-http transport)")
+    command: str | None = Field(default=None,
+                                description="Command to run for stdio transport (e.g. 'python' or 'docker')")
+    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: str | 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."""
+        if self.transport == "stdio":
+            if self.url is not None:
+                raise ValueError("url should not be set when using stdio transport")
+            if not self.command:
+                raise ValueError("command is required when using stdio transport")
+            # 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 streamable-http transport")
+            if not self.url:
+                raise ValueError("url is required when using streamable-http transport")
+
+        return self
+
+
+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_call_timeout: timedelta = Field(
+        default=timedelta(seconds=60),
+        description="Timeout (in seconds) for the MCP tool call. Defaults to 60 seconds.")
+    auth_flow_timeout: timedelta = Field(
+        default=timedelta(seconds=300),
+        description="Timeout (in seconds) for the MCP auth flow. When the tool call requires interactive \
+        authentication, this timeout is used. Defaults to 300 seconds.")
+    reconnect_enabled: bool = Field(
+        default=True,
+        description="Whether to enable reconnecting to the MCP server if the connection is lost. \
+        Defaults to True.")
+    reconnect_max_attempts: int = Field(default=2,
+                                        ge=0,
+                                        description="Maximum number of reconnect attempts. Defaults to 2.")
+    reconnect_initial_backoff: float = Field(
+        default=0.5, ge=0.0, description="Initial backoff time for reconnect attempts. Defaults to 0.5 seconds.")
+    reconnect_max_backoff: float = Field(
+        default=50.0, ge=0.0, description="Maximum backoff time for reconnect attempts. Defaults to 50 seconds.")
+    tool_overrides: dict[str, MCPToolOverrideConfig] | None = Field(
+        default=None,
+        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
+        """)
+    session_aware_tools: bool = Field(default=True,
+                                      description="Session-aware tools are created if True. Defaults to True.")
+    max_sessions: int = Field(default=100,
+                              ge=1,
+                              description="Maximum number of concurrent session clients. Defaults to 100.")
+    session_idle_timeout: timedelta = Field(
+        default=timedelta(hours=1),
+        description="Time after which inactive sessions are cleaned up. Defaults to 1 hour.")
+
+    @model_validator(mode="after")
+    def _validate_reconnect_backoff(self) -> "MCPClientConfig":
+        """Validate reconnect backoff values."""
+        if self.reconnect_max_backoff < self.reconnect_initial_backoff:
+            raise ValueError("reconnect_max_backoff must be greater than or equal to reconnect_initial_backoff")
+        return self