nvidia-nat-mcp 1.3.0a20250926__py3-none-any.whl → 1.3.0a20250928__py3-none-any.whl

nat/plugins/mcp/auth/auth_flow_handler.py

@@ -0,0 +1,144 @@
+# 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 asyncio
+import logging
+import secrets
+import webbrowser
+
+import pkce
+from authlib.integrations.httpx_client import AsyncOAuth2Client
+from fastapi import FastAPI
+
+from nat.authentication.oauth2.oauth2_auth_code_flow_provider_config import OAuth2AuthCodeFlowProviderConfig
+from nat.data_models.authentication import AuthenticatedContext
+from nat.data_models.authentication import AuthFlowType
+from nat.data_models.authentication import AuthProviderBaseConfig
+from nat.front_ends.console.authentication_flow_handler import ConsoleAuthenticationFlowHandler
+from nat.front_ends.console.authentication_flow_handler import _FlowState
+from nat.front_ends.fastapi.fastapi_front_end_controller import _FastApiFrontEndController
+
+logger = logging.getLogger(__name__)
+
+
+class MCPAuthenticationFlowHandler(ConsoleAuthenticationFlowHandler):
+    """
+    Authentication helper for MCP environments.
+
+    This handler is specifically designed for MCP tool discovery scenarios where
+    authentication needs to happen before the default auth_callback is available
+    in the Context. It handles OAuth2 authorization code flow during MCP client
+    startup and tool discovery phases.
+
+    Key differences from console handler:
+    - Only supports OAuth2 Authorization Code flow (no HTTP Basic)
+    - Optimized for MCP tool discovery workflows
+    - Designed for single-use authentication during startup
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._server_controller: _FastApiFrontEndController | None = None
+        self._redirect_app: FastAPI | None = None
+        self._server_lock = asyncio.Lock()
+        self._oauth_client: AsyncOAuth2Client | None = None
+
+    async def authenticate(self, config: AuthProviderBaseConfig, method: AuthFlowType) -> AuthenticatedContext:
+        """
+        Handle the OAuth2 authorization code flow for MCP environments.
+
+        Args:
+            config: OAuth2 configuration for MCP server
+            method: Authentication method (only OAUTH2_AUTHORIZATION_CODE supported)
+
+        Returns:
+            AuthenticatedContext with Bearer token for MCP server access
+
+        Raises:
+            ValueError: If config is invalid for MCP use case
+            NotImplementedError: If method is not OAuth2 Authorization Code
+        """
+        logger.info("Starting MCP authentication flow")
+
+        if method == AuthFlowType.OAUTH2_AUTHORIZATION_CODE:
+            if not isinstance(config, OAuth2AuthCodeFlowProviderConfig):
+                raise ValueError("Requested OAuth2 Authorization Code Flow but passed invalid config")
+
+            # MCP-specific validation
+            if not config.redirect_uri:
+                raise ValueError("MCP authentication requires redirect_uri to be configured")
+
+            logger.info("MCP authentication configured for server: %s", getattr(config, 'server_url', 'unknown'))
+            return await self._handle_oauth2_auth_code_flow(config)
+
+        raise NotImplementedError(f'Auth method "{method}" not supported for MCP environments')
+
+    async def _handle_oauth2_auth_code_flow(self, cfg: OAuth2AuthCodeFlowProviderConfig) -> AuthenticatedContext:
+        logger.info("Starting MCP OAuth2 authorization code flow")
+
+        state = secrets.token_urlsafe(16)
+        flow_state = _FlowState()
+        client = self.construct_oauth_client(cfg)
+
+        flow_state.token_url = cfg.token_url
+        flow_state.use_pkce = cfg.use_pkce
+
+        # PKCE bits
+        if cfg.use_pkce:
+            verifier, challenge = pkce.generate_pkce_pair()
+            flow_state.verifier = verifier
+            flow_state.challenge = challenge
+            logger.debug("PKCE enabled for MCP authentication")
+
+        auth_url, _ = client.create_authorization_url(
+            cfg.authorization_url,
+            state=state,
+            code_verifier=flow_state.verifier if cfg.use_pkce else None,
+            code_challenge=flow_state.challenge if cfg.use_pkce else None,
+            **(cfg.authorization_kwargs or {})
+        )
+
+        async with self._server_lock:
+            if self._redirect_app is None:
+                self._redirect_app = await self._build_redirect_app()
+
+            await self._start_redirect_server()
+            self._flows[state] = flow_state
+
+        logger.info("MCP authentication: Your browser has been opened for authentication.")
+        logger.info("This will authenticate you with the MCP server for tool discovery.")
+        webbrowser.open(auth_url)
+
+        # Use default timeout for MCP tool discovery
+        timeout = 300
+
+        try:
+            token = await asyncio.wait_for(flow_state.future, timeout=timeout)
+            logger.info("MCP authentication successful, token obtained")
+        except asyncio.TimeoutError as exc:
+            logger.error("MCP authentication timed out")
+            raise RuntimeError(f"MCP authentication timed out ({timeout} seconds). Please try again.") from exc
+        finally:
+            async with self._server_lock:
+                self._flows.pop(state, None)
+                await self._stop_redirect_server()
+
+        return AuthenticatedContext(
+            headers={"Authorization": f"Bearer {token['access_token']}"},
+            metadata={
+                "expires_at": token.get("expires_at"),
+                "raw_token": token,
+            },
+        )

nat/plugins/mcp/auth/auth_provider.py

@@ -14,6 +14,8 @@
 # limitations under the License.
 
 import logging
+from collections.abc import Awaitable
+from typing import Callable
 from urllib.parse import urljoin
 from urllib.parse import urlparse
 
@@ -26,10 +28,12 @@ from mcp.shared.auth import OAuthClientInformationFull
 from mcp.shared.auth import OAuthClientMetadata
 from mcp.shared.auth import OAuthMetadata
 from mcp.shared.auth import ProtectedResourceMetadata
+from nat.authentication.interfaces import AuthenticatedContext
+from nat.authentication.interfaces import AuthFlowType
 from nat.authentication.interfaces import AuthProviderBase
-from nat.data_models.authentication import AuthReason
-from nat.data_models.authentication import AuthRequest
+from nat.authentication.oauth2.oauth2_auth_code_flow_provider_config import OAuth2AuthCodeFlowProviderConfig
 from nat.data_models.authentication import AuthResult
+from nat.plugins.mcp.auth.auth_flow_handler import MCPAuthenticationFlowHandler
 from nat.plugins.mcp.auth.auth_provider_config import MCPOAuth2ProviderConfig
 
 logger = logging.getLogger(__name__)
@@ -40,6 +44,7 @@ class OAuth2Endpoints(BaseModel):
     authorization_url: HttpUrl = Field(..., description="OAuth2 authorization endpoint URL")
     token_url: HttpUrl = Field(..., description="OAuth2 token endpoint URL")
     registration_url: HttpUrl | None = Field(default=None, description="OAuth2 client registration endpoint URL")
+    scopes: list[str] | None = Field(default=None, description="OAuth2 scopes to be used for the authentication")
 
 
 class OAuth2Credentials(BaseModel):
@@ -60,9 +65,11 @@ class DiscoverOAuth2Endpoints:
     def __init__(self, config: MCPOAuth2ProviderConfig):
         self.config = config
         self._cached_endpoints: OAuth2Endpoints | None = None
-        self._last_oauth_scopes: list[str] | None = None
+        self._authenticated_servers: dict[str, AuthResult] = {}
 
-    async def discover(self, reason: AuthReason, www_authenticate: str | None) -> tuple[OAuth2Endpoints, bool]:
+        self._flow_handler: MCPAuthenticationFlowHandler = MCPAuthenticationFlowHandler()
+
+    async def discover(self, response: httpx.Response | None = None) -> tuple[OAuth2Endpoints, bool]:
         """
         Discover OAuth2 endpoints from MCP server.
 
@@ -73,21 +80,24 @@ class DiscoverOAuth2Endpoints:
         Returns:
             A tuple of OAuth2Endpoints and a boolean indicating if the endpoints have changed.
         """
+        is_401_retry = response is not None and response.status_code == 401
         # Fast path: reuse cache when not a 401 retry
-        if reason != AuthReason.RETRY_AFTER_401 and self._cached_endpoints is not None:
+        if not is_401_retry and self._cached_endpoints is not None:
             return self._cached_endpoints, False
 
         issuer: str = str(self.config.server_url)  # default to server URL
         endpoints: OAuth2Endpoints | None = None
 
         # 1) 401 hint (RFC 9728) if present
-        if reason == AuthReason.RETRY_AFTER_401 and www_authenticate:
-            hint_url = self._extract_from_www_authenticate_header(www_authenticate)
-            if hint_url:
-                logger.info("Using RFC 9728 resource_metadata hint: %s", hint_url)
-                issuer_hint = await self._fetch_pr_issuer(hint_url)
-                if issuer_hint:
-                    issuer = issuer_hint
+        if is_401_retry and response:
+            www_authenticate = response.headers.get("WWW-Authenticate")
+            if www_authenticate:
+                hint_url = self._extract_from_www_authenticate_header(www_authenticate)
+                if hint_url:
+                    logger.info("Using RFC 9728 resource_metadata hint: %s", hint_url)
+                    issuer_hint = await self._fetch_pr_issuer(hint_url)
+                    if issuer_hint:
+                        issuer = issuer_hint
 
         # 2) Try RS protected resource well-known if we still only have default issuer
         if issuer == str(self.config.server_url):
@@ -105,10 +115,7 @@ class DiscoverOAuth2Endpoints:
         if endpoints is None:
             raise RuntimeError("Could not discover OAuth2 endpoints from MCP server")
 
-        changed = (self._cached_endpoints is None
-                   or endpoints.authorization_url != self._cached_endpoints.authorization_url
-                   or endpoints.token_url != self._cached_endpoints.token_url
-                   or endpoints.registration_url != self._cached_endpoints.registration_url)
+        changed = (self._cached_endpoints is None or endpoints.model_dump() != self._cached_endpoints.model_dump())
         self._cached_endpoints = endpoints
         logger.info("OAuth2 endpoints selected: %s", self._cached_endpoints)
         return self._cached_endpoints, changed
@@ -155,10 +162,29 @@ class DiscoverOAuth2Endpoints:
         async with httpx.AsyncClient(timeout=10.0) as client:
             for url in urls:
                 try:
-                    resp = await client.get(url, headers={"Accept": "application/json"})
+                    resp = await client.get(url, follow_redirects=True, headers={"Accept": "application/json"})
                     if resp.status_code != 200:
                         continue
+
+                    # Check content type before attempting JSON parsing
+                    content_type = resp.headers.get("content-type", "").lower()
+                    if "application/json" not in content_type:
+                        logger.info(
+                            "Discovery endpoint %s returned non-JSON content type: %s. "
+                            "This may indicate the endpoint doesn't support discovery or requires authentication.",
+                            url,
+                            content_type)
+                        # If it's HTML, log a more helpful message
+                        if "text/html" in content_type:
+                            logger.info("The endpoint appears to be returning an HTML page instead of OAuth metadata. "
+                                        "This often means:")
+                            logger.info("1. The OAuth discovery endpoint doesn't exist at this URL")
+                            logger.info("2. The server requires authentication before providing discovery metadata")
+                            logger.info("3. The URL is pointing to a web application instead of an OAuth server")
+                        continue
+
                     body = await resp.aread()
+
                     try:
                         meta = OAuthMetadata.model_validate_json(body)
                     except Exception as e:
@@ -167,14 +193,18 @@ class DiscoverOAuth2Endpoints:
                     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
-                        self._last_oauth_scopes = meta.scopes_supported
                         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,
+                            scopes=meta.scopes_supported,
                         )
                 except Exception as e:
                     logger.debug("Discovery failed at %s: %s", url, e)
+
+        # If we get here, all discovery URLs failed
+        logger.info("OAuth discovery failed for all attempted URLs.")
+        logger.info("Attempted URLs: %s", urls)
         return None
 
     def _build_path_aware_discovery_urls(self, base_or_issuer: str) -> list[str]:
@@ -184,17 +214,19 @@ class DiscoverOAuth2Endpoints:
         path = (p.path or "").rstrip("/")
         urls: list[str] = []
         if path:
-            urls.append(urljoin(base, f"/.well-known/oauth-authorization-server{path}"))
+            # this is the specified by the MCP spec
+            urls.append(urljoin(base, f".well-known/oauth-protected-resource{path}"))
+            # this is fallback for backward compatibility
+            urls.append(urljoin(base, f"{path}/.well-known/oauth-authorization-server"))
         urls.append(urljoin(base, "/.well-known/oauth-authorization-server"))
         if path:
-            urls.append(urljoin(base, f"/.well-known/openid-configuration{path}"))
+            # this is the specified by the MCP spec
+            urls.append(urljoin(base, f".well-known/openid-configuration{path}"))
+            # this is fallback for backward compatibility
+            urls.append(urljoin(base, f"{path}/.well-known/openid-configuration"))
         urls.append(base_or_issuer.rstrip("/") + "/.well-known/openid-configuration")
         return urls
 
-    def scopes_supported(self) -> list[str] | None:
-        """Get the last OAuth scopes discovered from the AS."""
-        return self._last_oauth_scopes
-
 
 class DynamicClientRegistration:
     """Dynamic client registration utility."""
@@ -264,51 +296,54 @@ class MCPOAuth2Provider(AuthProviderBase[MCPOAuth2ProviderConfig]):
 
         # For the OAuth2 flow
         self._auth_code_provider = None
+        self._flow_handler = MCPAuthenticationFlowHandler()
 
-    async def authenticate(self, user_id: str | None = None) -> AuthResult:
+        self._auth_callback = None
+
+    def _set_custom_auth_callback(self,
+                                  auth_callback: Callable[[OAuth2AuthCodeFlowProviderConfig, AuthFlowType],
+                                                          Awaitable[AuthenticatedContext]]):
+        """Set the custom authentication callback."""
+        if not self._auth_callback:
+            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)
+
+    async def authenticate(self, user_id: str | None = None, **kwargs) -> AuthResult:
         """
         Authenticate using MCP OAuth2 flow via NAT framework.
+
+        If response is provided in kwargs (typically from a 401), performs:
         1. Dynamic endpoints discovery (RFC9728 + RFC 8414 + OIDC)
         2. Client registration (RFC7591)
-        3. Use NAT's standard OAuth2 flow (OAuth2AuthCodeFlowProvider)
+        3. Authentication
+
+        Otherwise, performs standard authentication flow.
         """
-        auth_request = self.config.auth_request
-        if not auth_request:
-            auth_request = AuthRequest(reason=AuthReason.NORMAL)
-
-        if auth_request.reason != AuthReason.RETRY_AFTER_401:
-            # auth provider is expected to be setup via 401, till that time we return empty auth result
-            if not self._auth_code_provider:
-                return AuthResult(credentials=[], token_expires_at=None, raw={})
-
-        await self._discover_and_register(auth_request)
-        # Use NAT's standard OAuth2 flow
-        if auth_request.reason == AuthReason.RETRY_AFTER_401:
-            # force fresh delegate (clears in-mem token cache)
-            self._auth_code_provider = None
-            # preserve other fields, just normalize reason & inject user_id
-            auth_request = auth_request.model_copy(update={
-                "reason": AuthReason.NORMAL, "user_id": user_id, "www_authenticate": None
-            })
-        # back-compat: propagate user_id if provided but not set in the request
-        elif user_id is not None and auth_request.user_id is None:
-            auth_request = auth_request.model_copy(update={"user_id": user_id})
-
-        # Perform the OAuth2 flow without lock
-        return await self._perform_oauth2_flow(auth_request=auth_request)
-
-    async def _discover_and_register(self, auth_request: AuthRequest):
+        response = kwargs.get('response')
+        if response and response.status_code == 401:
+            await self._discover_and_register(response=response)
+
+        return await self._nat_oauth2_authenticate(user_id=user_id)
+
+    @property
+    def _effective_scopes(self) -> list[str]:
+        """Get the effective scopes to be used for the authentication."""
+        return self.config.scopes or (self._cached_endpoints.scopes if self._cached_endpoints else []) or []
+
+    async def _discover_and_register(self, response: httpx.Response | None = None):
         """
         Discover OAuth2 endpoints and register an OAuth2 client with the Authorization Server
         using OIDC client registration.
         """
         # Discover OAuth2 endpoints
-        self._cached_endpoints, endpoints_changed = await self._discoverer.discover(reason=auth_request.reason,
-                                                                                    www_authenticate=auth_request.www_authenticate)
+        self._cached_endpoints, endpoints_changed = await self._discoverer.discover(response=response)
         if endpoints_changed:
             logger.info("OAuth2 endpoints: %s", self._cached_endpoints)
             self._cached_credentials = None  # invalidate credentials tied to old AS
-        effective_scopes = self._effective_scopes()
+            self._auth_code_provider = None
+        effective_scopes = self._effective_scopes
 
         # Client registration
         if not self._cached_credentials:
@@ -324,21 +359,20 @@ class MCPOAuth2Provider(AuthProviderBase[MCPOAuth2ProviderConfig]):
                 self._cached_credentials = await self._registrar.register(self._cached_endpoints, effective_scopes)
                 logger.info("Registered OAuth2 client: %s", self._cached_credentials.client_id)
 
-    def _effective_scopes(self) -> list[str] | None:
-        """
-        Prefer caller-provided scopes; otherwise fall back to AS-advertised scopes_supported.
-        """
-        return self.config.scopes or self._discoverer.scopes_supported()
-
-    async def _build_oauth2_delegate(self):
-        """Build NAT OAuth2 provider and delegate auth token acquisition and refresh to it"""
+    async def _nat_oauth2_authenticate(self, user_id: str | None = None) -> AuthResult:
+        """Perform the OAuth2 flow using MCP-specific authentication flow handler."""
         from nat.authentication.oauth2.oauth2_auth_code_flow_provider import OAuth2AuthCodeFlowProvider
-        from nat.authentication.oauth2.oauth2_auth_code_flow_provider_config import OAuth2AuthCodeFlowProviderConfig
+
+        if not self._cached_endpoints or not self._cached_credentials:
+            # if discovery is yet to to be done return empty auth result
+            return AuthResult(credentials=[], token_expires_at=None, raw={})
 
         endpoints = self._cached_endpoints
         credentials = self._cached_credentials
 
+        # Build the OAuth2 provider if not already built
         if self._auth_code_provider is None:
+            scopes = self._effective_scopes
             oauth2_config = OAuth2AuthCodeFlowProviderConfig(
                 client_id=credentials.client_id,
                 client_secret=credentials.client_secret or "",
@@ -346,22 +380,15 @@ class MCPOAuth2Provider(AuthProviderBase[MCPOAuth2ProviderConfig]):
                 token_url=str(endpoints.token_url),
                 token_endpoint_auth_method=getattr(self.config, "token_endpoint_auth_method", None),
                 redirect_uri=str(self.config.redirect_uri) if self.config.redirect_uri else "",
-                scopes=self._effective_scopes() or [],
+                scopes=scopes,
                 use_pkce=bool(self.config.use_pkce),
-            )
-
+                authorization_kwargs={"resource": str(self.config.server_url)})
             self._auth_code_provider = OAuth2AuthCodeFlowProvider(oauth2_config)
 
-    async def _perform_oauth2_flow(self, auth_request: AuthRequest | None = None) -> AuthResult:
-        """Perform the OAuth2 flow using NAT OAuth2 provider."""
-        # This helper is only for non-401 flows
-        if auth_request and auth_request.reason == AuthReason.RETRY_AFTER_401:
-            raise RuntimeError("_perform_oauth2_flow should not be called for RETRY_AFTER_401")
-
-        if not self._cached_endpoints or not self._cached_credentials:
-            raise RuntimeError("OAuth2 flow called before discovery/registration")
+            # 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)
 
-        # (Re)build the delegate if needed
-        await self._build_oauth2_delegate()
-        # Let the delegate handle per-user cache + refresh
-        return await self._auth_code_provider.authenticate()
+        # 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

@@ -18,7 +18,6 @@ from pydantic import HttpUrl
 from pydantic import model_validator
 
 from nat.authentication.interfaces import AuthProviderBaseConfig
-from nat.data_models.authentication import AuthRequest
 
 
 class MCPOAuth2ProviderConfig(AuthProviderBaseConfig, name="mcp_oauth2"):
@@ -51,12 +50,16 @@ class MCPOAuth2ProviderConfig(AuthProviderBaseConfig, name="mcp_oauth2"):
     # Advanced options
     use_pkce: bool = Field(default=True, description="Use PKCE for authorization code flow")
 
-    auth_request: AuthRequest | None = Field(default=None, description="Auth request for authentication (metadata)")
+    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")
 
     @model_validator(mode="after")
     def validate_auth_config(self):
         """Validate authentication configuration for MCP-specific options."""
 
+        # if default_user_id is not provided, use the server_url as the default user id
+        if not self.default_user_id:
+            self.default_user_id = str(self.server_url)
         # Dynamic registration + MCP discovery
         if self.enable_dynamic_registration and not self.client_id:
             # Pure dynamic registration - no explicit credentials needed

nat/plugins/mcp/client_base.py

@@ -21,10 +21,12 @@ import logging
 from abc import ABC
 from abc import abstractmethod
 from collections.abc import AsyncGenerator
+from collections.abc import Callable
 from contextlib import AsyncExitStack
 from contextlib import asynccontextmanager
 from datetime import timedelta
 
+import anyio
 import httpx
 
 from mcp import ClientSession
@@ -33,9 +35,9 @@ 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 AuthenticatedContext
+from nat.authentication.interfaces import AuthFlowType
 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 convert_to_mcp_error
 from nat.plugins.mcp.exception_handler import format_mcp_error
 from nat.plugins.mcp.exception_handler import mcp_exception_handler
@@ -53,74 +55,89 @@ class AuthAdapter(httpx.Auth):
     Converts AuthProviderBase to httpx.Auth interface for dynamic token management.
     """
 
-    def __init__(self, auth_provider: AuthProviderBase, auth_for_tool_calls_only: bool = False):
+    def __init__(self, auth_provider: AuthProviderBase):
         self.auth_provider = auth_provider
-        self.auth_for_tool_calls_only = auth_for_tool_calls_only
+        # each adapter instance has its own lock to avoid unnecessary delays for multiple clients
+        self._lock = anyio.Lock()
 
     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:
+        async with self._lock:
             try:
-                # Get fresh auth headers with 401 context
-                auth_headers = await self._get_auth_headers(reason=AuthReason.RETRY_AFTER_401, response=response)
+                # Get auth headers from the NAT auth provider:
+                # 1. If discovery is yet to done this will return None and request will be sent without auth header.
+                # 2. If discovery is done, this will return the auth header from cache if the token is still valid
+                auth_headers = await self._get_auth_headers(request=request, response=None)
                 request.headers.update(auth_headers)
-                yield request  # Retry the request
             except Exception as e:
-                logger.info("Failed to refresh auth after 401: %s", 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:
+                    # 401 can happen if:
+                    # 1. The request was sent without auth header
+                    # 2. The auth headers are invalid
+                    # 3. The auth headers are expired
+                    # 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
+                    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)
         return
 
-    def _is_tool_call_request(self, request: httpx.Request) -> bool:
-        """Check if this is a tool call request based on the request body."""
+    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"):
-                    return True
+                    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 False
+        return None, False
 
-    async def _get_auth_headers(self, reason: AuthReason, response: httpx.Response | None = None) -> dict[str, str]:
+    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."""
-        # 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()
+            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 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
+            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)
+
             # 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:
-                logger.warning("Auth provider did not return BearerTokenCred")
+                logger.info("Auth provider did not return BearerTokenCred")
                 return {}
         except Exception as e:
             logger.warning("Failed to get auth token: %s", e)
@@ -155,6 +172,7 @@ class MCPBaseClient(ABC):
         self._initial_connection = False
 
         # Convert auth provider to AuthAdapter
+        self._auth_provider = auth_provider
         self._httpx_auth = AuthAdapter(auth_provider) if auth_provider else None
 
         self._tool_call_timeout = tool_call_timeout
@@ -314,6 +332,34 @@ class MCPBaseClient(ABC):
             raise MCPToolNotFoundError(tool_name, self.server_name)
         return tool
 
+    def set_user_auth_callback(self, auth_callback: Callable[[AuthFlowType], AuthenticatedContext]):
+        """Set the user authentication callback."""
+        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):
 
@@ -539,9 +585,35 @@ class MCPToolClient:
         Args:
             tool_args (dict[str, Any]): A dictionary of key value pairs to serve as inputs for the MCP tool.
         """
-        logger.info("Calling tool %s with arguments %s", self._tool_name, tool_args)
+        if self._session is None:
+            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")
+        except Exception:
+            pass
+
         try:
-            result = await self._parent_client.call_tool(self._tool_name, tool_args)
+            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)
 
             output = []
             for res in result.content:

nat/plugins/mcp/client_impl.py

@@ -56,7 +56,8 @@ class MCPServerConfig(BaseModel):
     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")
+    auth_provider: str | AuthenticationRef | None = Field(default=None,
+                                                          description="Reference to authentication provider")
 
     @model_validator(mode="after")
     def validate_model(self):
@@ -92,7 +93,8 @@ class MCPClientConfig(FunctionGroupBaseConfig, name="mcp_client"):
     """
     server: MCPServerConfig = Field(..., description="Server connection details (transport, url/command, etc.)")
     tool_call_timeout: timedelta = Field(
-        default=timedelta(seconds=5), description="Timeout (in seconds) for the MCP tool call. Defaults to 5 seconds.")
+        default=timedelta(seconds=60),
+        description="Timeout (in seconds) for the MCP tool call. Defaults to 60 seconds.")
     reconnect_enabled: bool = Field(
         default=True,
         description="Whether to enable reconnecting to the MCP server if the connection is lost. \

{nvidia_nat_mcp-1.3.0a20250926.dist-info → nvidia_nat_mcp-1.3.0a20250928.dist-info}/METADATA

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

{nvidia_nat_mcp-1.3.0a20250926.dist-info → nvidia_nat_mcp-1.3.0a20250928.dist-info}/RECORD

@@ -1,18 +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=uOPn1EE7iL6yvHeBFAdQlNoNmtCaBplX5oyu7AeO2lQ,21356
-nat/plugins/mcp/client_impl.py,sha256=5QdA6nt3xQcA_g-YsGZ5BwBsBVoIlk9ObI4mdJNiuYU,10685
+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_provider.py,sha256=GOmM9vCfVd0QiyD_hBj7zCfkiimHa1WDfTTWWfMsr_k,17466
-nat/plugins/mcp/auth/auth_provider_config.py,sha256=bE6IKV0yveo98KXr0xynrH5BMwPhRv8xbaMBwYu42YQ,3587
+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.0a20250926.dist-info/METADATA,sha256=Ax4EGQZ3KdGWPm9-MyZ6x6DPm5-6SBPmuPYCFWXD2Sw,1997
-nvidia_nat_mcp-1.3.0a20250926.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nvidia_nat_mcp-1.3.0a20250926.dist-info/entry_points.txt,sha256=rYvUp4i-klBr3bVNh7zYOPXret704vTjvCk1qd7FooI,97
-nvidia_nat_mcp-1.3.0a20250926.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
-nvidia_nat_mcp-1.3.0a20250926.dist-info/RECORD,,
+nvidia_nat_mcp-1.3.0a20250928.dist-info/METADATA,sha256=VnuxYtll39Hir1_hDQvhqRWPM00M6A9JD3NCd7JISaw,1997
+nvidia_nat_mcp-1.3.0a20250928.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_mcp-1.3.0a20250928.dist-info/entry_points.txt,sha256=rYvUp4i-klBr3bVNh7zYOPXret704vTjvCk1qd7FooI,97
+nvidia_nat_mcp-1.3.0a20250928.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
+nvidia_nat_mcp-1.3.0a20250928.dist-info/RECORD,,