nvidia-nat-mcp 1.3.0a20250930__py3-none-any.whl → 1.3.0a20251002__py3-none-any.whl

nat/plugins/mcp/client_base.py

@@ -59,6 +59,8 @@ class AuthAdapter(httpx.Auth):
         self.auth_provider = auth_provider
         # 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,11 +87,21 @@ 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]:
@@ -148,12 +160,19 @@ 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),
+                 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,
@@ -173,6 +192,7 @@ class MCPBaseClient(ABC):
         self._httpx_auth = AuthAdapter(auth_provider) 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 +287,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 +314,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 +387,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
         }
 
@@ -361,12 +434,7 @@ class MCPBaseClient(ABC):
         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
+            timeout = await self._get_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)
@@ -376,7 +444,8 @@ class MCPBaseClient(ABC):
 
         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 +460,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 +511,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,7 +570,8 @@ class MCPStreamableHTTPClient(MCPBaseClient):
     def __init__(self,
                  url: str,
                  auth_provider: AuthProviderBase | 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,
@@ -505,6 +579,7 @@ class MCPStreamableHTTPClient(MCPBaseClient):
         super().__init__("streamable-http",
                          auth_provider=auth_provider,
                          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 +625,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.")
@@ -643,7 +716,7 @@ class MCPToolClient:
                 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)
+                result = await self._parent_client.call_tool(self._tool_name, tool_args)
 
             output = []
             for res in result.content:

nat/plugins/mcp/client_impl.py

@@ -139,6 +139,10 @@ class MCPClientConfig(FunctionGroupBaseConfig, name="mcp_client"):
     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. \
@@ -196,7 +200,8 @@ async def mcp_client_function_group(config: MCPClientConfig, _builder: Builder):
         client = MCPStdioClient(config.server.command,
                                 config.server.args,
                                 config.server.env,
-                                config.tool_call_timeout,
+                                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,
@@ -204,6 +209,7 @@ async def mcp_client_function_group(config: MCPClientConfig, _builder: Builder):
     elif 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,
@@ -212,6 +218,7 @@ async def mcp_client_function_group(config: MCPClientConfig, _builder: Builder):
         client = MCPStreamableHTTPClient(str(config.server.url),
                                          auth_provider=auth_provider,
                                          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,

{nvidia_nat_mcp-1.3.0a20250930.dist-info → nvidia_nat_mcp-1.3.0a20251002.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nvidia-nat-mcp
-Version: 1.3.0a20250930
+Version: 1.3.0a20251002
 Summary: Subpackage for MCP client integration in NeMo Agent toolkit
 Keywords: ai,rag,agents,mcp
 Classifier: Programming Language :: Python
@@ -9,7 +9,7 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Requires-Python: <3.14,>=3.11
 Description-Content-Type: text/markdown
-Requires-Dist: nvidia-nat==v1.3.0a20250930
+Requires-Dist: nvidia-nat==v1.3.0a20251002
 Requires-Dist: mcp~=1.14
 
 <!--

{nvidia_nat_mcp-1.3.0a20250930.dist-info → nvidia_nat_mcp-1.3.0a20251002.dist-info}/RECORD

@@ -1,7 +1,7 @@
 nat/meta/pypi.md,sha256=GyV4DI1d9ThgEhnYTQ0vh40Q9hPC8jN-goLnRiFDmZ8,1498
 nat/plugins/mcp/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
-nat/plugins/mcp/client_base.py,sha256=UPdQ4vxe4MXVTF05jJx20dn8JPTvQN1ugZ-dups3-44,26418
-nat/plugins/mcp/client_impl.py,sha256=sjSFIJeD6LkexN5IbDq_OuoKW52ulpp26LpwaUjOpwg,13142
+nat/plugins/mcp/client_base.py,sha256=x_mgrCORXqYfJcYZg8zd0wm1AN1uFl51A7bYye0X-rc,30151
+nat/plugins/mcp/client_impl.py,sha256=FGWlpzyBDR2tNkV9Ek7brSX9EWh98kfAGCr49zMzDSU,13657
 nat/plugins/mcp/exception_handler.py,sha256=4JVdZDJL4LyumZEcMIEBK2LYC6djuSMzqUhQDZZ6dUo,7648
 nat/plugins/mcp/exceptions.py,sha256=EGVOnYlui8xufm8dhJyPL1SUqBLnCGOTvRoeyNcmcWE,5980
 nat/plugins/mcp/register.py,sha256=HOT2Wl2isGuyFc7BUTi58-BbjI5-EtZMZo7stsv5pN4,831
@@ -12,8 +12,8 @@ nat/plugins/mcp/auth/auth_flow_handler.py,sha256=2JgK0aH-5ouQCd2ov0lDMJAD5ZWIQJ7
 nat/plugins/mcp/auth/auth_provider.py,sha256=OfxPCEaXuhP8anOdrTRH-_E78CrbJtzW6i81_kebpDk,19321
 nat/plugins/mcp/auth/auth_provider_config.py,sha256=vhU47Vcp_30M8tWu0FumbJ6pdUnFbBZm-ABdNlup__U,3821
 nat/plugins/mcp/auth/register.py,sha256=yzphsn1I4a5G39_IacbuX0ZQqGM8fevvTUM_B94UXKE,1211
-nvidia_nat_mcp-1.3.0a20250930.dist-info/METADATA,sha256=CBJk-2cOfsLjvFNIieoKpxaMHYtAeRFBWfzvj1Ti1Ns,1997
-nvidia_nat_mcp-1.3.0a20250930.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nvidia_nat_mcp-1.3.0a20250930.dist-info/entry_points.txt,sha256=rYvUp4i-klBr3bVNh7zYOPXret704vTjvCk1qd7FooI,97
-nvidia_nat_mcp-1.3.0a20250930.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
-nvidia_nat_mcp-1.3.0a20250930.dist-info/RECORD,,
+nvidia_nat_mcp-1.3.0a20251002.dist-info/METADATA,sha256=IxP-J_3VJWx9SXiCSgQOsrBdagj6i24bx_XK8KZeI-0,1997
+nvidia_nat_mcp-1.3.0a20251002.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_mcp-1.3.0a20251002.dist-info/entry_points.txt,sha256=rYvUp4i-klBr3bVNh7zYOPXret704vTjvCk1qd7FooI,97
+nvidia_nat_mcp-1.3.0a20251002.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
+nvidia_nat_mcp-1.3.0a20251002.dist-info/RECORD,,