nvidia-nat-mcp 1.3.0a20251006__py3-none-any.whl → 1.3.0a20251008__py3-none-any.whl

nat/meta/pypi.md

@@ -19,9 +19,9 @@ limitations under the License.
 
 
 # NVIDIA NeMo Agent Toolkit MCP Subpackage
-Subpackage for MCP client integration in NeMo Agent toolkit.
+Subpackage for MCP integration in NeMo Agent toolkit.
 
-This package provides MCP (Model Context Protocol) client functionality, allowing NeMo Agent toolkit workflows to connect to external MCP servers and use their tools as functions.
+This package provides MCP (Model Context Protocol) functionality, allowing NeMo Agent toolkit workflows to connect to external MCP servers and use their tools as functions.
 
 ## Features
 

nat/plugins/mcp/auth/auth_provider.py

@@ -23,6 +23,7 @@ import httpx
 from pydantic import BaseModel
 from pydantic import Field
 from pydantic import HttpUrl
+from pydantic import TypeAdapter
 
 from mcp.shared.auth import OAuthClientInformationFull
 from mcp.shared.auth import OAuthClientMetadata
@@ -65,7 +66,6 @@ class DiscoverOAuth2Endpoints:
     def __init__(self, config: MCPOAuth2ProviderConfig):
         self.config = config
         self._cached_endpoints: OAuth2Endpoints | None = None
-        self._authenticated_servers: dict[str, AuthResult] = {}
 
         self._flow_handler: MCPAuthenticationFlowHandler = MCPAuthenticationFlowHandler()
 
@@ -192,11 +192,13 @@ class DiscoverOAuth2Endpoints:
                         continue
                     if meta.authorization_endpoint and meta.token_endpoint:
                         logger.info("Discovered OAuth2 endpoints from %s", url)
-                        # this is bit of a hack to get the scopes supported by the auth server
+                        # Convert AnyHttpUrl to HttpUrl using TypeAdapter
+                        http_url_adapter = TypeAdapter(HttpUrl)
                         return OAuth2Endpoints(
-                            authorization_url=str(meta.authorization_endpoint),
-                            token_url=str(meta.token_endpoint),
-                            registration_url=str(meta.registration_endpoint) if meta.registration_endpoint else None,
+                            authorization_url=http_url_adapter.validate_python(str(meta.authorization_endpoint)),
+                            token_url=http_url_adapter.validate_python(str(meta.token_endpoint)),
+                            registration_url=http_url_adapter.validate_python(str(meta.registration_endpoint))
+                            if meta.registration_endpoint else None,
                             scopes=meta.scopes_supported,
                         )
                 except Exception as e:
@@ -283,8 +285,9 @@ class DynamicClientRegistration:
 class MCPOAuth2Provider(AuthProviderBase[MCPOAuth2ProviderConfig]):
     """MCP OAuth2 authentication provider that delegates to NAT framework."""
 
-    def __init__(self, config: MCPOAuth2ProviderConfig):
+    def __init__(self, config: MCPOAuth2ProviderConfig, builder=None):
         super().__init__(config)
+        self._builder = builder
 
         # Discovery
         self._discoverer = DiscoverOAuth2Endpoints(config)
@@ -300,6 +303,19 @@ class MCPOAuth2Provider(AuthProviderBase[MCPOAuth2ProviderConfig]):
 
         self._auth_callback = None
 
+        # Initialize token storage
+        self._token_storage = None
+        self._token_storage_object_store_name = None
+
+        if self.config.token_storage_object_store:
+            # Store object store name, will be resolved later when builder context is available
+            self._token_storage_object_store_name = self.config.token_storage_object_store
+            logger.info(f"Configured to use object store '{self._token_storage_object_store_name}' for token storage")
+        else:
+            # Default: use in-memory token storage
+            from .token_storage import InMemoryTokenStorage
+            self._token_storage = InMemoryTokenStorage()
+
     def _set_custom_auth_callback(self,
                                   auth_callback: Callable[[OAuth2AuthCodeFlowProviderConfig, AuthFlowType],
                                                           Awaitable[AuthenticatedContext]]):
@@ -308,7 +324,7 @@ class MCPOAuth2Provider(AuthProviderBase[MCPOAuth2ProviderConfig]):
             logger.info("Using custom authentication callback")
             self._auth_callback = auth_callback
             if self._auth_code_provider:
-                self._auth_code_provider._set_custom_auth_callback(self._auth_callback)
+                self._auth_code_provider._set_custom_auth_callback(self._auth_callback)  # type: ignore[arg-type]
 
     async def authenticate(self, user_id: str | None = None, **kwargs) -> AuthResult:
         """
@@ -374,6 +390,22 @@ class MCPOAuth2Provider(AuthProviderBase[MCPOAuth2ProviderConfig]):
         endpoints = self._cached_endpoints
         credentials = self._cached_credentials
 
+        # Resolve object store reference if needed
+        if self._token_storage_object_store_name and not self._token_storage:
+            try:
+                if not self._builder:
+                    raise RuntimeError("Builder not available for resolving object store")
+                object_store = await self._builder.get_object_store_client(self._token_storage_object_store_name)
+                from .token_storage import ObjectStoreTokenStorage
+                self._token_storage = ObjectStoreTokenStorage(object_store)
+                logger.info(f"Initialized token storage with object store '{self._token_storage_object_store_name}'")
+            except Exception as e:
+                logger.warning(
+                    f"Failed to resolve object store '{self._token_storage_object_store_name}' for token storage: {e}. "
+                    "Falling back to in-memory storage.")
+                from .token_storage import InMemoryTokenStorage
+                self._token_storage = InMemoryTokenStorage()
+
         # Build the OAuth2 provider if not already built
         if self._auth_code_provider is None:
             scopes = self._effective_scopes
@@ -387,12 +419,12 @@ class MCPOAuth2Provider(AuthProviderBase[MCPOAuth2ProviderConfig]):
                 scopes=scopes,
                 use_pkce=bool(self.config.use_pkce),
                 authorization_kwargs={"resource": str(self.config.server_url)})
-            self._auth_code_provider = OAuth2AuthCodeFlowProvider(oauth2_config)
+            self._auth_code_provider = OAuth2AuthCodeFlowProvider(oauth2_config, token_storage=self._token_storage)
 
             # Use MCP-specific authentication method if available
             if hasattr(self._auth_code_provider, "_set_custom_auth_callback"):
-                self._auth_code_provider._set_custom_auth_callback(self._auth_callback
-                                                                   or self._flow_handler.authenticate)
+                callback = self._auth_callback or self._flow_handler.authenticate
+                self._auth_code_provider._set_custom_auth_callback(callback)  # type: ignore[arg-type]
 
         # Auth code provider is responsible for per-user cache + refresh
         return await self._auth_code_provider.authenticate(user_id=user_id)

nat/plugins/mcp/auth/auth_provider_config.py

@@ -53,6 +53,11 @@ class MCPOAuth2ProviderConfig(AuthProviderBaseConfig, name="mcp_oauth2"):
     default_user_id: str | None = Field(default=None, description="Default user ID for authentication")
     allow_default_user_id_for_tool_calls: bool = Field(default=True, description="Allow default user ID for tool calls")
 
+    # Token storage configuration
+    token_storage_object_store: str | None = Field(
+        default=None,
+        description="Reference to object store for secure token storage. If None, uses in-memory storage.")
+
     @model_validator(mode="after")
     def validate_auth_config(self):
         """Validate authentication configuration for MCP-specific options."""

nat/plugins/mcp/auth/register.py

@@ -22,4 +22,4 @@ from nat.plugins.mcp.auth.auth_provider_config import MCPOAuth2ProviderConfig
 @register_auth_provider(config_type=MCPOAuth2ProviderConfig)
 async def mcp_oauth2_provider(authentication_provider: MCPOAuth2ProviderConfig, builder: Builder):
     """Register MCP OAuth2 authentication provider with NAT system."""
-    yield MCPOAuth2Provider(authentication_provider)
+    yield MCPOAuth2Provider(authentication_provider, builder=builder)

nat/plugins/mcp/auth/token_storage.py

@@ -0,0 +1,265 @@
+# 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.
+
+import hashlib
+import json
+import logging
+from abc import ABC
+from abc import abstractmethod
+
+from nat.data_models.authentication import AuthResult
+from nat.data_models.authentication import BasicAuthCred
+from nat.data_models.authentication import BearerTokenCred
+from nat.data_models.authentication import CookieCred
+from nat.data_models.authentication import HeaderCred
+from nat.data_models.authentication import QueryCred
+from nat.data_models.object_store import NoSuchKeyError
+from nat.object_store.interfaces import ObjectStore
+from nat.object_store.models import ObjectStoreItem
+
+logger = logging.getLogger(__name__)
+
+
+class TokenStorageBase(ABC):
+    """
+    Abstract base class for token storage implementations.
+
+    Token storage implementations handle the secure persistence of authentication
+    tokens for MCP OAuth2 flows. Implementations can use various backends such as
+    object stores, databases, or in-memory storage.
+    """
+
+    @abstractmethod
+    async def store(self, user_id: str, auth_result: AuthResult) -> None:
+        """
+        Store an authentication result for a user.
+
+        Args:
+            user_id: The unique identifier for the user
+            auth_result: The authentication result to store
+        """
+        pass
+
+    @abstractmethod
+    async def retrieve(self, user_id: str) -> AuthResult | None:
+        """
+        Retrieve an authentication result for a user.
+
+        Args:
+            user_id: The unique identifier for the user
+
+        Returns:
+            The authentication result if found, None otherwise
+        """
+        pass
+
+    @abstractmethod
+    async def delete(self, user_id: str) -> None:
+        """
+        Delete an authentication result for a user.
+
+        Args:
+            user_id: The unique identifier for the user
+        """
+        pass
+
+    @abstractmethod
+    async def clear_all(self) -> None:
+        """
+        Clear all stored authentication results.
+        """
+        pass
+
+
+class ObjectStoreTokenStorage(TokenStorageBase):
+    """
+    Token storage implementation backed by a NeMo Agent toolkit object store.
+
+    This implementation uses the object store infrastructure to persist tokens,
+    which provides encryption at rest, access controls, and persistence across
+    restarts when using backends like S3, MySQL, or Redis.
+    """
+
+    def __init__(self, object_store: ObjectStore):
+        """
+        Initialize the object store token storage.
+
+        Args:
+            object_store: The object store instance to use for token persistence
+        """
+        self._object_store = object_store
+
+    def _get_key(self, user_id: str) -> str:
+        """
+        Generate the object store key for a user's token.
+
+        Uses SHA256 hash to ensure the key is S3-compatible and doesn't
+        contain special characters like "://" that are invalid in object keys.
+
+        Args:
+            user_id: The user identifier
+
+        Returns:
+            The object store key
+        """
+        # Hash the user_id to create an S3-safe key
+        user_hash = hashlib.sha256(user_id.encode('utf-8')).hexdigest()
+        return f"tokens/{user_hash}"
+
+    async def store(self, user_id: str, auth_result: AuthResult) -> None:
+        """
+        Store an authentication result in the object store.
+
+        Args:
+            user_id: The unique identifier for the user
+            auth_result: The authentication result to store
+        """
+        key = self._get_key(user_id)
+
+        # Serialize the AuthResult to JSON with secrets exposed
+        # SecretStr values are masked by default, so we need to expose them manually
+        # Create a serializable dict with exposed secrets
+        auth_dict = auth_result.model_dump(mode='json')
+        # Manually expose SecretStr values in credentials
+        for i, cred_obj in enumerate(auth_result.credentials):
+            if isinstance(cred_obj, BearerTokenCred):
+                auth_dict['credentials'][i]['token'] = cred_obj.token.get_secret_value()
+            elif isinstance(cred_obj, BasicAuthCred):
+                auth_dict['credentials'][i]['username'] = cred_obj.username.get_secret_value()
+                auth_dict['credentials'][i]['password'] = cred_obj.password.get_secret_value()
+            elif isinstance(cred_obj, HeaderCred | QueryCred | CookieCred):
+                auth_dict['credentials'][i]['value'] = cred_obj.value.get_secret_value()
+
+        data = json.dumps(auth_dict).encode('utf-8')
+
+        # Prepare metadata
+        metadata = {}
+        if auth_result.token_expires_at:
+            metadata["expires_at"] = auth_result.token_expires_at.isoformat()
+
+        # Create the object store item
+        item = ObjectStoreItem(data=data, content_type="application/json", metadata=metadata if metadata else None)
+
+        # Store using upsert to handle both new and existing tokens
+        await self._object_store.upsert_object(key, item)
+
+    async def retrieve(self, user_id: str) -> AuthResult | None:
+        """
+        Retrieve an authentication result from the object store.
+
+        Args:
+            user_id: The unique identifier for the user
+
+        Returns:
+            The authentication result if found, None otherwise
+        """
+        key = self._get_key(user_id)
+
+        try:
+            item = await self._object_store.get_object(key)
+            # Deserialize the AuthResult from JSON
+            auth_result = AuthResult.model_validate_json(item.data)
+            return auth_result
+        except NoSuchKeyError:
+            return None
+        except Exception as e:
+            logger.error(f"Error deserializing token for user {user_id}: {e}", exc_info=True)
+            return None
+
+    async def delete(self, user_id: str) -> None:
+        """
+        Delete an authentication result from the object store.
+
+        Args:
+            user_id: The unique identifier for the user
+        """
+        key = self._get_key(user_id)
+
+        try:
+            await self._object_store.delete_object(key)
+        except NoSuchKeyError:
+            # Token doesn't exist, which is fine for delete operations
+            pass
+
+    async def clear_all(self) -> None:
+        """
+        Clear all stored authentication results.
+
+        Note: This implementation does not support clearing all tokens as the
+        object store interface doesn't provide a list operation. Individual
+        tokens must be deleted explicitly.
+        """
+        logger.warning("clear_all() is not supported for ObjectStoreTokenStorage")
+
+
+class InMemoryTokenStorage(TokenStorageBase):
+    """
+    In-memory token storage using NeMo Agent toolkit's built-in object store.
+
+    This implementation uses the in-memory object store for token persistence,
+    which provides a secure default option that doesn't require external storage
+    configuration. Tokens are stored in memory and cleared when the process exits.
+    """
+
+    def __init__(self):
+        """
+        Initialize the in-memory token storage.
+        """
+        from nat.object_store.in_memory_object_store import InMemoryObjectStore
+
+        # Create a dedicated in-memory object store for tokens
+        self._object_store = InMemoryObjectStore()
+
+        # Wrap with ObjectStoreTokenStorage for the actual implementation
+        self._storage = ObjectStoreTokenStorage(self._object_store)
+        logger.debug("Initialized in-memory token storage")
+
+    async def store(self, user_id: str, auth_result: AuthResult) -> None:
+        """
+        Store an authentication result in memory.
+
+        Args:
+            user_id: The unique identifier for the user
+            auth_result: The authentication result to store
+        """
+        await self._storage.store(user_id, auth_result)
+
+    async def retrieve(self, user_id: str) -> AuthResult | None:
+        """
+        Retrieve an authentication result from memory.
+
+        Args:
+            user_id: The unique identifier for the user
+
+        Returns:
+            The authentication result if found, None otherwise
+        """
+        return await self._storage.retrieve(user_id)
+
+    async def delete(self, user_id: str) -> None:
+        """
+        Delete an authentication result from memory.
+
+        Args:
+            user_id: The unique identifier for the user
+        """
+        await self._storage.delete(user_id)
+
+    async def clear_all(self) -> None:
+        """
+        Clear all stored authentication results from memory.
+        """
+        # For in-memory storage, we can access the internal storage
+        self._object_store._store.clear()

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,8 +54,9 @@ 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
@@ -104,41 +104,13 @@ class AuthAdapter(httpx.Auth):
                     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
@@ -171,6 +143,7 @@ class MCPBaseClient(ABC):
     def __init__(self,
                  transport: str = 'streamable-http',
                  auth_provider: AuthProviderBase | None = None,
+                 user_id: str | None = None,
                  tool_call_timeout: timedelta = timedelta(seconds=60),
                  auth_flow_timeout: timedelta = timedelta(seconds=300),
                  reconnect_enabled: bool = True,
@@ -189,7 +162,9 @@ 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
@@ -421,24 +396,6 @@ 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))
-            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)
-
     @mcp_exception_handler
     async def call_tool(self, tool_name: str, tool_args: dict | None):
 
@@ -570,6 +527,7 @@ class MCPStreamableHTTPClient(MCPBaseClient):
     def __init__(self,
                  url: str,
                  auth_provider: AuthProviderBase | None = None,
+                 user_id: str | None = None,
                  tool_call_timeout: timedelta = timedelta(seconds=60),
                  auth_flow_timeout: timedelta = timedelta(seconds=300),
                  reconnect_enabled: bool = True,
@@ -578,6 +536,7 @@ class MCPStreamableHTTPClient(MCPBaseClient):
                  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,
@@ -662,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.
@@ -698,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._parent_client.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: