nvidia-nat-mcp 1.3.0a20250929__py3-none-any.whl → 1.3.0rc1__py3-none-any.whl

nat/plugins/mcp/auth/auth_flow_handler.py

@@ -127,7 +127,7 @@ class MCPAuthenticationFlowHandler(ConsoleAuthenticationFlowHandler):
         try:
             token = await asyncio.wait_for(flow_state.future, timeout=timeout)
             logger.info("MCP authentication successful, token obtained")
-        except asyncio.TimeoutError as exc:
+        except TimeoutError as exc:
             logger.error("MCP authentication timed out")
             raise RuntimeError(f"MCP authentication timed out ({timeout} seconds). Please try again.") from exc
         finally:

nat/plugins/mcp/auth/auth_provider.py

@@ -15,7 +15,7 @@
 
 import logging
 from collections.abc import Awaitable
-from typing import Callable
+from collections.abc import Callable
 from urllib.parse import urljoin
 from urllib.parse import urlparse
 
@@ -321,6 +321,10 @@ class MCPOAuth2Provider(AuthProviderBase[MCPOAuth2ProviderConfig]):
 
         Otherwise, performs standard authentication flow.
         """
+        if not user_id:
+            # MCP tool calls cannot be made without an authorized user
+            raise RuntimeError("User is not authorized to call the tool")
+
         response = kwargs.get('response')
         if response and response.status_code == 401:
             await self._discover_and_register(response=response)

nat/plugins/mcp/client_base.py

@@ -120,11 +120,8 @@ class AuthAdapter(httpx.Auth):
                 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 if it exists, default user id can be used if allowed
-                if self.auth_provider.config.allow_default_user_id_for_tool_calls:
-                    user_id = session_id or self.auth_provider.config.default_user_id
-                else:
-                    user_id = session_id
+                # 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
@@ -184,6 +181,10 @@ class MCPBaseClient(ABC):
         self._reconnect_max_backoff = reconnect_max_backoff
         self._reconnect_lock: asyncio.Lock = asyncio.Lock()
 
+    @property
+    def auth_provider(self) -> AuthProviderBase | None:
+        return self._auth_provider
+
     @property
     def transport(self) -> str:
         return self._transport
@@ -280,7 +281,8 @@ class MCPBaseClient(ABC):
                 return await coro()
             raise
 
-    async def get_tools(self):
+    @mcp_exception_handler
+    async def get_tools(self) -> dict[str, MCPToolClient]:
         """
         Retrieve a dictionary of all tools served by the MCP server.
         Uses unauthenticated session for discovery.
@@ -288,7 +290,16 @@ class MCPBaseClient(ABC):
 
         async def _get_tools():
             session = self._session
-            return await session.list_tools()
+            try:
+                # Add timeout to the list_tools call.
+                # This is needed because MCP SDK does not support timeout for list_tools()
+                with anyio.fail_after(self._tool_call_timeout.total_seconds()):
+                    tools = await session.list_tools()
+            except TimeoutError as e:
+                from nat.plugins.mcp.exceptions import MCPTimeoutError
+                raise MCPTimeoutError(self.server_name, e)
+
+            return tools
 
         try:
             response = await self._with_reconnect(_get_tools)
@@ -536,7 +547,7 @@ class MCPToolClient:
 
     def __init__(self,
                  session: ClientSession,
-                 parent_client: "MCPBaseClient",
+                 parent_client: MCPBaseClient,
                  tool_name: str,
                  tool_description: str | None,
                  tool_input_schema: dict | None = None,
@@ -578,6 +589,32 @@ 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.
@@ -589,25 +626,18 @@ class MCPToolClient:
             raise RuntimeError("No session available for tool call")
 
         # Extract context information
-        session_id = None
         try:
-            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
-            cookies = getattr(_Ctx.get().metadata, "cookies", None)
-            if cookies:
-                session_id = cookies.get("nat-session")
+            session_id = self._get_session_id()
         except Exception:
-            pass
+            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)
@@ -630,6 +660,6 @@ class MCPToolClient:
 
         except MCPError as e:
             format_mcp_error(e, include_traceback=False)
-            result_str = "MCPToolClient tool call failed: %s" % e.original_exception
+            result_str = f"MCPToolClient tool call failed: {e.original_exception}"
 
         return result_str

nat/plugins/mcp/client_impl.py

@@ -32,6 +32,50 @@ from nat.plugins.mcp.tool import mcp_tool_function
 logger = logging.getLogger(__name__)
 
 
+class MCPFunctionGroup(FunctionGroup):
+    """
+    A specialized FunctionGroup for MCP clients that includes MCP-specific attributes
+    with proper type safety.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # MCP client attributes with proper typing
+        self._mcp_client = None  # Will be set to the actual MCP client instance
+        self._mcp_client_server_name: str | None = None
+        self._mcp_client_transport: str | None = None
+
+    @property
+    def mcp_client(self):
+        """Get the MCP client instance."""
+        return self._mcp_client
+
+    @mcp_client.setter
+    def mcp_client(self, client):
+        """Set the MCP client instance."""
+        self._mcp_client = client
+
+    @property
+    def mcp_client_server_name(self) -> str | None:
+        """Get the MCP client server name."""
+        return self._mcp_client_server_name
+
+    @mcp_client_server_name.setter
+    def mcp_client_server_name(self, server_name: str | None):
+        """Set the MCP client server name."""
+        self._mcp_client_server_name = server_name
+
+    @property
+    def mcp_client_transport(self) -> str | None:
+        """Get the MCP client transport type."""
+        return self._mcp_client_transport
+
+    @mcp_client_transport.setter
+    def mcp_client_transport(self, transport: str | None):
+        """Set the MCP client transport type."""
+        self._mcp_client_transport = transport
+
+
 class MCPToolOverrideConfig(BaseModel):
     """
     Configuration for overriding tool properties when exposing from MCP server.
@@ -177,10 +221,16 @@ async def mcp_client_function_group(config: MCPClientConfig, _builder: Builder):
 
     logger.info("Configured to use MCP server at %s", client.server_name)
 
-    # Create the function group
-    group = FunctionGroup(config=config)
+    # Create the MCP function group
+    group = MCPFunctionGroup(config=config)
 
     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)
 
@@ -196,12 +246,24 @@ async def mcp_client_function_group(config: MCPClientConfig, _builder: Builder):
             # Create the tool function
             tool_fn = mcp_tool_function(tool)
 
+            # Normalize optional typing for linter/type-checker compatibility
+            single_fn = tool_fn.single_fn
+            if single_fn is None:
+                # Should not happen because mcp_tool_function 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.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,
+                               fn=single_fn,
+                               input_schema=input_schema,
                                converters=tool_fn.converters)
 
         yield group

nat/plugins/mcp/exception_handler.py

@@ -94,7 +94,7 @@ def extract_primary_exception(exceptions: list[Exception]) -> Exception:
     """
     # Prioritize connection errors
     for exc in exceptions:
-        if isinstance(exc, (httpx.ConnectError, ConnectionError)):
+        if isinstance(exc, httpx.ConnectError | ConnectionError):
             return exc
 
     # Then timeout errors

{nvidia_nat_mcp-1.3.0a20250929.dist-info → nvidia_nat_mcp-1.3.0rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nvidia-nat-mcp
-Version: 1.3.0a20250929
+Version: 1.3.0rc1
 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.0a20250929
+Requires-Dist: nvidia-nat==v1.3.0-rc1
 Requires-Dist: mcp~=1.14
 
 <!--

nvidia_nat_mcp-1.3.0rc1.dist-info/RECORD

@@ -0,0 +1,19 @@
+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/exception_handler.py,sha256=4JVdZDJL4LyumZEcMIEBK2LYC6djuSMzqUhQDZZ6dUo,7648
+nat/plugins/mcp/exceptions.py,sha256=EGVOnYlui8xufm8dhJyPL1SUqBLnCGOTvRoeyNcmcWE,5980
+nat/plugins/mcp/register.py,sha256=HOT2Wl2isGuyFc7BUTi58-BbjI5-EtZMZo7stsv5pN4,831
+nat/plugins/mcp/tool.py,sha256=v3MFsiaLJy8Ourcfqa6ohtAE2Nn-vqpC6Q6gsCdJ28Q,6165
+nat/plugins/mcp/utils.py,sha256=3fuzYpC14wrfMOTOGvY2KHWcxZvBWqrxdDZD17lhmC8,4055
+nat/plugins/mcp/auth/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
+nat/plugins/mcp/auth/auth_flow_handler.py,sha256=2JgK0aH-5ouQCd2ov0lDMJAD5ZWIQJ7SVcXaLArxn6Y,6010
+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.0rc1.dist-info/METADATA,sha256=NdfOVgW-bDqAZWInj84TQgDjH05tlBwEYO-mZdlwklM,1986
+nvidia_nat_mcp-1.3.0rc1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_mcp-1.3.0rc1.dist-info/entry_points.txt,sha256=rYvUp4i-klBr3bVNh7zYOPXret704vTjvCk1qd7FooI,97
+nvidia_nat_mcp-1.3.0rc1.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
+nvidia_nat_mcp-1.3.0rc1.dist-info/RECORD,,

nvidia_nat_mcp-1.3.0a20250929.dist-info/RECORD

@@ -1,19 +0,0 @@
-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=x6NHs3X2alyB6Wo4pjWZrGyuU7lLqGC_cLO9fuL4Zgw,25194
-nat/plugins/mcp/client_impl.py,sha256=M1gTMlp3RLhFaAHOvwkk38boFy05MixV_glrIEcMjvo,10759
-nat/plugins/mcp/exception_handler.py,sha256=JdPdZG1NgWpdRnIz7JTGHiJASS5wot9nJiD3SRWV4Kw,7649
-nat/plugins/mcp/exceptions.py,sha256=EGVOnYlui8xufm8dhJyPL1SUqBLnCGOTvRoeyNcmcWE,5980
-nat/plugins/mcp/register.py,sha256=HOT2Wl2isGuyFc7BUTi58-BbjI5-EtZMZo7stsv5pN4,831
-nat/plugins/mcp/tool.py,sha256=v3MFsiaLJy8Ourcfqa6ohtAE2Nn-vqpC6Q6gsCdJ28Q,6165
-nat/plugins/mcp/utils.py,sha256=3fuzYpC14wrfMOTOGvY2KHWcxZvBWqrxdDZD17lhmC8,4055
-nat/plugins/mcp/auth/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
-nat/plugins/mcp/auth/auth_flow_handler.py,sha256=eRqRS2t-YzJ3kEFaG0PEC8DzctYzaJzr9XLZGkvuxq0,6018
-nat/plugins/mcp/auth/auth_provider.py,sha256=5TTaPlIXMgwrE4YcZ_HO9-GNBBFaDTdUKAa5fuALayI,19142
-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.0a20250929.dist-info/METADATA,sha256=HceYMMAdbXe5IBiiXJIyrd5IdEUcWIHgjJtiW53BT4I,1997
-nvidia_nat_mcp-1.3.0a20250929.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nvidia_nat_mcp-1.3.0a20250929.dist-info/entry_points.txt,sha256=rYvUp4i-klBr3bVNh7zYOPXret704vTjvCk1qd7FooI,97
-nvidia_nat_mcp-1.3.0a20250929.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
-nvidia_nat_mcp-1.3.0a20250929.dist-info/RECORD,,