nvidia-nat-mcp 1.3.0rc6__py3-none-any.whl → 1.4.0a20251123__py3-none-any.whl

nat/plugins/mcp/auth/auth_provider_config.py

@@ -18,6 +18,7 @@ from pydantic import HttpUrl
 from pydantic import model_validator
 
 from nat.authentication.interfaces import AuthProviderBaseConfig
+from nat.data_models.common import OptionalSecretStr
 
 
 class MCPOAuth2ProviderConfig(AuthProviderBaseConfig, name="mcp_oauth2"):
@@ -36,7 +37,8 @@ class MCPOAuth2ProviderConfig(AuthProviderBaseConfig, name="mcp_oauth2"):
 
     # Client registration (manual registration vs DCR)
     client_id: str | None = Field(default=None, description="OAuth2 client ID for pre-registered clients")
-    client_secret: str | None = Field(default=None, description="OAuth2 client secret for pre-registered clients")
+    client_secret: OptionalSecretStr = Field(default=None,
+                                             description="OAuth2 client secret for pre-registered clients")
     enable_dynamic_registration: bool = Field(default=True,
                                               description="Enable OAuth2 Dynamic Client Registration (RFC 7591)")
     client_name: str = Field(default="NAT MCP Client", description="OAuth2 client name for dynamic registration")

nat/plugins/mcp/auth/register.py

@@ -17,9 +17,17 @@ from nat.builder.builder import Builder
 from nat.cli.register_workflow import register_auth_provider
 from nat.plugins.mcp.auth.auth_provider import MCPOAuth2Provider
 from nat.plugins.mcp.auth.auth_provider_config import MCPOAuth2ProviderConfig
+from nat.plugins.mcp.auth.service_account.provider import MCPServiceAccountProvider
+from nat.plugins.mcp.auth.service_account.provider_config import MCPServiceAccountProviderConfig
 
 
 @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, builder=builder)
+
+
+@register_auth_provider(config_type=MCPServiceAccountProviderConfig)
+async def mcp_service_account_provider(authentication_provider: MCPServiceAccountProviderConfig, builder: Builder):
+    """Register MCP Service Account authentication provider with NAT system."""
+    yield MCPServiceAccountProvider(authentication_provider, builder=builder)

nat/plugins/mcp/auth/service_account/__init__.py

@@ -0,0 +1,14 @@
+# 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.

nat/plugins/mcp/auth/service_account/provider.py

@@ -0,0 +1,136 @@
+# 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 importlib
+import logging
+import typing
+
+from pydantic import SecretStr
+
+from nat.authentication.interfaces import AuthProviderBase
+from nat.data_models.authentication import AuthResult
+from nat.data_models.authentication import Credential
+from nat.data_models.authentication import HeaderCred
+from nat.plugins.mcp.auth.service_account.provider_config import MCPServiceAccountProviderConfig
+from nat.plugins.mcp.auth.service_account.token_client import ServiceAccountTokenClient
+
+logger = logging.getLogger(__name__)
+
+
+class MCPServiceAccountProvider(AuthProviderBase[MCPServiceAccountProviderConfig]):
+    """
+    MCP service account authentication provider using OAuth2 client credentials.
+
+    Provides headless authentication for MCP clients using service account credentials.
+    Supports two authentication patterns:
+
+    1. Single authentication: OAuth2 service account token only
+    2. Dual authentication: OAuth2 service account token + service-specific token
+
+    """
+
+    def __init__(self, config: MCPServiceAccountProviderConfig, builder=None):
+        super().__init__(config)
+
+        # Initialize token client
+        self._token_client = ServiceAccountTokenClient(
+            client_id=config.client_id,
+            client_secret=config.client_secret,
+            token_url=config.token_url,
+            scopes=" ".join(config.scopes),  # Convert list to space-delimited string for OAuth2
+            token_cache_buffer_seconds=config.token_cache_buffer_seconds,
+        )
+
+        # Load dynamic service token function if configured
+        self._service_token_function = None
+        if config.service_token and config.service_token.function:
+            self._service_token_function = self._load_function(config.service_token.function)
+
+        logger.info("Initialized MCP service account auth provider: "
+                    "token_url=%s, scopes=%s, has_service_token=%s",
+                    config.token_url,
+                    config.scopes,
+                    config.service_token is not None)
+
+    def _load_function(self, function_path: str) -> typing.Callable:
+        """Load a Python function from a module path string (e.g., 'my_module.get_token')."""
+        try:
+            module_name, func_name = function_path.rsplit(".", 1)
+            module = importlib.import_module(module_name)
+            func = getattr(module, func_name)
+            logger.info("Loaded service token function: %s", function_path)
+            return func
+        except Exception as e:
+            raise ValueError(f"Failed to load service token function '{function_path}': {e}") from e
+
+    async def authenticate(self, user_id: str | None = None, **kwargs) -> AuthResult:
+        """
+        Authenticate using OAuth2 client credentials flow.
+
+        Note: user_id is ignored for service accounts (non-session-specific).
+
+        Returns:
+            AuthResult with HeaderCred objects for service account authentication
+        """
+        # Get OAuth2 access token (cached if still valid)
+        access_token = await self._token_client.get_access_token()
+
+        # Build credentials list using HeaderCred
+        credentials: list[Credential] = [
+            HeaderCred(name="Authorization", value=SecretStr(f"Bearer {access_token.get_secret_value()}"))
+        ]
+
+        # Add service-specific token if configured
+        if self.config.service_token:
+            service_header = self.config.service_token.header
+            service_token_value = None
+
+            # Get service token from static config or dynamic function
+            if self.config.service_token.token:
+                # Static token from config
+                service_token_value = self.config.service_token.token.get_secret_value()
+
+            elif self._service_token_function:
+                # Dynamic token from function
+                try:
+                    # Pass configured kwargs to the function
+                    # Function can access runtime context via AIQContext.get() if needed
+                    # Handle both sync and async functions
+                    if asyncio.iscoroutinefunction(self._service_token_function):
+                        result = await self._service_token_function(**self.config.service_token.kwargs)
+                    else:
+                        result = self._service_token_function(**self.config.service_token.kwargs)
+
+                    # Handle function return type: str or tuple[str, str]
+                    if isinstance(result, tuple):
+                        service_header, service_token_value = result
+                    else:
+                        service_token_value = result
+
+                    logger.debug("Retrieved service token via dynamic function")
+
+                except Exception as e:
+                    raise RuntimeError(f"Failed to get service token from function: {e}") from e
+
+            if service_token_value:
+                credentials.append(HeaderCred(name=service_header, value=SecretStr(service_token_value)))
+
+        # Return AuthResult with HeaderCred objects
+        return AuthResult(
+            credentials=credentials,
+            token_expires_at=self._token_client.token_expires_at,
+            raw={},
+        )

nat/plugins/mcp/auth/service_account/provider_config.py

@@ -0,0 +1,137 @@
+# 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 typing
+
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic import field_validator
+from pydantic import model_validator
+
+from nat.authentication.interfaces import AuthProviderBaseConfig
+from nat.data_models.common import OptionalSecretStr
+from nat.data_models.common import SerializableSecretStr
+
+
+class ServiceTokenConfig(BaseModel):
+    """
+    Configuration for service-specific token in dual authentication patterns.
+
+    Supports two modes:
+
+    1. Static token: Provide token and header directly
+    2. Dynamic function: Provide function path and optional kwargs
+
+    The function will be called on every request and should have signature::
+
+        async def get_service_token(**kwargs) -> str | tuple[str, str]
+
+    If function returns ``tuple[str, str]``, it's interpreted as (header_name, token).
+    If function returns ``str``, it's the token and header field is used for header name.
+
+    The function can access runtime context via AIQContext.get() if needed.
+    """
+
+    # Static token approach
+    token: OptionalSecretStr = Field(
+        default=None,
+        description="Static service token value (mutually exclusive with function)",
+    )
+
+    header: str = Field(
+        default="X-Service-Account-Token",
+        description="HTTP header name for service token (default: 'X-Service-Account-Token')",
+    )
+
+    # Dynamic function approach
+    function: str | None = Field(
+        default=None,
+        description=("Python function path that returns service token dynamically (mutually exclusive with token). "
+                     "Function signature: async def func(\\**kwargs) -> str | tuple[str, str]. "
+                     "Access runtime context via AIQContext.get() if needed."),
+    )
+
+    kwargs: dict[str, typing.Any] = Field(
+        default_factory=dict,
+        description="Additional keyword arguments to pass to the custom function",
+    )
+
+    @model_validator(mode="after")
+    def validate_token_or_function(self):
+        """Ensure either token or function is provided, but not both."""
+        has_token = self.token is not None
+        has_function = self.function is not None
+
+        if not has_token and not has_function:
+            raise ValueError("Either 'token' or 'function' must be provided in service_token config")
+
+        if has_token and has_function:
+            raise ValueError("Cannot specify both 'token' and 'function' in service_token config. Choose one.")
+
+        return self
+
+
+class MCPServiceAccountProviderConfig(AuthProviderBaseConfig, name="mcp_service_account"):
+    """
+    Configuration for MCP service account authentication using OAuth2 client credentials.
+
+    Generic implementation supporting any OAuth2 client credentials flow.
+
+    Supports two authentication patterns:
+    1. Single authentication: OAuth2 service account token only
+    2. Dual authentication: OAuth2 service account token + service-specific token
+
+    Common use cases:
+    - Headless/automated MCP workflows
+    - CI/CD pipelines
+    - Backend services without user interaction
+
+    All values must be provided via configuration. Use ${ENV_VAR} syntax in YAML
+    configs for environment variable substitution.
+    """
+
+    # Required: OAuth2 client credentials
+    client_id: str = Field(description="OAuth2 client identifier")
+
+    client_secret: SerializableSecretStr = Field(description="OAuth2 client secret")
+
+    # Required: Token endpoint URL
+    token_url: str = Field(description="OAuth2 token endpoint URL")
+
+    # Required: OAuth2 scopes
+    scopes: list[str] = Field(description="List of OAuth2 scopes (will be joined with spaces for OAuth2 request)")
+
+    # Optional: Service-specific token configuration for dual authentication patterns
+    service_token: ServiceTokenConfig | None = Field(
+        default=None,
+        description="Optional service token configuration for dual authentication patterns. "
+        "Provide either a static token or a dynamic function that returns the token at runtime.",
+    )
+
+    # Token caching configuration
+    token_cache_buffer_seconds: int = Field(default=300,
+                                            description="Seconds before token expiry to refresh (default: 300s/5min)")
+
+    @field_validator("scopes", mode="before")
+    @classmethod
+    def validate_scopes(cls, v):
+        """
+        Accept both list[str] and space-delimited string formats for scopes.
+        Converts string to list for consistency.
+        """
+        if isinstance(v, str):
+            # Split space-delimited string into list
+            return [scope.strip() for scope in v.split() if scope.strip()]
+        return v

nat/plugins/mcp/auth/service_account/token_client.py

@@ -0,0 +1,156 @@
+# 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 base64
+import logging
+from datetime import datetime
+from datetime import timedelta
+
+import httpx
+from pydantic import SecretStr
+
+logger = logging.getLogger(__name__)
+
+
+class ServiceAccountTokenClient:
+    """
+    Generic OAuth2 client credentials token client for service accounts.
+
+    Implements standard OAuth2 client credentials flow with token caching.
+    """
+
+    def __init__(
+        self,
+        client_id: str,
+        client_secret: SecretStr,
+        token_url: str,
+        scopes: str,
+        token_cache_buffer_seconds: int = 300,
+    ):
+        """
+        Initialize service account token client.
+
+        Args:
+            client_id: OAuth2 client identifier
+            client_secret: OAuth2 client secret (SecretStr)
+            token_url: OAuth2 token endpoint URL
+            scopes: Space-separated list of scopes
+            token_cache_buffer_seconds: Seconds before expiry to refresh (default: 5 min)
+        """
+        self.client_id = client_id
+        self.client_secret = client_secret
+        self.token_url = token_url
+        self.scopes = scopes
+        self.token_cache_buffer_seconds = token_cache_buffer_seconds
+
+        # Token cache
+        self._cached_token: SecretStr | None = None
+        self._token_expires_at: datetime | None = None
+        self._lock = None  # Will be initialized as asyncio.Lock when needed
+
+    @property
+    def token_expires_at(self) -> datetime | None:
+        return self._token_expires_at
+
+    async def _get_lock(self) -> asyncio.Lock:
+        """Lazy initialization of asyncio.Lock."""
+        if self._lock is None:
+            self._lock = asyncio.Lock()
+        return self._lock
+
+    def _is_token_valid(self) -> bool:
+        """Check if cached token is still valid (with buffer time)."""
+        if not self._cached_token or not self._token_expires_at:
+            return False
+        buffer = timedelta(seconds=self.token_cache_buffer_seconds)
+        return datetime.now() < (self._token_expires_at - buffer)
+
+    async def get_access_token(self) -> SecretStr:
+        """
+        Get OAuth2 access token, using cache if valid.
+
+        Returns:
+            Access token as SecretStr
+
+        Raises:
+            RuntimeError: If token acquisition fails
+        """
+        # Fast path: check cache without lock
+        if self._is_token_valid():
+            logger.debug("Using cached service account token")
+            assert self._cached_token is not None  # _is_token_valid() ensures this
+            return self._cached_token
+
+        # Slow path: acquire lock and refresh token
+        lock = await self._get_lock()
+        async with lock:
+            # Double-check after acquiring lock
+            if self._is_token_valid():
+                logger.debug("Using cached service account token (acquired during lock wait)")
+                assert self._cached_token is not None  # _is_token_valid() ensures this
+                return self._cached_token
+
+            logger.info("Fetching new service account token")
+            return await self._fetch_new_token()
+
+    async def _fetch_new_token(self) -> SecretStr:
+        """
+        Fetch a new token from the OAuth2 token endpoint.
+
+        Returns:
+            New access token as SecretStr
+
+        Raises:
+            RuntimeError: If token request fails
+        """
+        # Encode credentials for Basic authentication
+        credentials = f"{self.client_id}:{self.client_secret.get_secret_value()}"
+        encoded_credentials = base64.b64encode(credentials.encode()).decode()
+
+        headers = {"Authorization": f"Basic {encoded_credentials}", "Content-Type": "application/x-www-form-urlencoded"}
+
+        data = {"grant_type": "client_credentials", "scope": self.scopes}
+
+        try:
+            async with httpx.AsyncClient(timeout=30.0) as client:
+                response = await client.post(self.token_url, headers=headers, data=data)
+
+                if response.status_code == 200:
+                    token_data = response.json()
+
+                    # Cache the token
+                    access_token = token_data.get("access_token")
+                    if not access_token:
+                        raise RuntimeError("Access token not found in token response")
+                    self._cached_token = SecretStr(access_token)
+                    expires_in = token_data.get("expires_in", 3600)
+                    self._token_expires_at = datetime.now() + timedelta(seconds=expires_in)
+
+                    logger.info("Service account token acquired (expires in %ss)", expires_in)
+                    return self._cached_token
+
+                elif response.status_code == 401:
+                    raise RuntimeError("Invalid service account credentials")
+                elif response.status_code == 429:
+                    raise RuntimeError("Service account rate limit exceeded")
+                else:
+                    raise RuntimeError(
+                        f"Service account token request failed: {response.status_code} - {response.text}")
+
+        except httpx.TimeoutException as e:
+            raise RuntimeError(f"Service account token request timed out: {e}") from e
+        except httpx.RequestError as e:
+            raise RuntimeError(f"Service account token request failed: {e}") from e

nat/plugins/mcp/client_base.py

@@ -112,14 +112,21 @@ class AuthAdapter(httpx.Auth):
             # 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
+            # Build headers from credentials
             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.info("Auth provider did not return BearerTokenCred")
-                return {}
+            from nat.data_models.authentication import HeaderCred
+            headers = {}
+
+            for cred in auth_result.credentials:
+                if isinstance(cred, BearerTokenCred):
+                    # Standard Bearer token
+                    token = cred.token.get_secret_value()
+                    headers["Authorization"] = f"Bearer {token}"
+                elif isinstance(cred, HeaderCred):
+                    # Generic header credential (supports custom formats and service accounts)
+                    headers[cred.name] = cred.value.get_secret_value()
+
+            return headers
         except Exception as e:
             logger.warning("Failed to get auth token: %s", e)
             return {}

nat/plugins/mcp/client_impl.py

@@ -109,6 +109,10 @@ class MCPFunctionGroup(FunctionGroup):
         self._shared_auth_provider: AuthProviderBase | None = None
         self._client_config: MCPClientConfig | None = None
 
+        # Auth provider config defaults (set when auth provider is assigned)
+        self._default_user_id: str | None = None
+        self._allow_default_user_id_for_tool_calls: bool = True
+
         # Use random session id for testing only
         self._use_random_session_id_for_testing: bool = False
 
@@ -176,9 +180,8 @@ class MCPFunctionGroup(FunctionGroup):
 
             if not session_id:
                 # use default user id if allowed
-                if self._shared_auth_provider and \
-                    self._shared_auth_provider.config.allow_default_user_id_for_tool_calls:
-                    session_id = self._shared_auth_provider.config.default_user_id
+                if self._shared_auth_provider and self._allow_default_user_id_for_tool_calls:
+                    session_id = self._default_user_id
             return session_id
         except Exception:
             return None
@@ -266,8 +269,7 @@ class MCPFunctionGroup(FunctionGroup):
         # If the session_id equals the configured default_user_id use the base client
         # instead of creating a per-session client
         if self._shared_auth_provider:
-            default_uid = self._shared_auth_provider.config.default_user_id
-            if default_uid and session_id == default_uid:
+            if self._default_user_id and session_id == self._default_user_id:
                 return self.mcp_client
 
         # Fast path: check if session already exists (reader lock for concurrent access)
@@ -435,8 +437,7 @@ def mcp_session_tool_function(tool, function_group: MCPFunctionGroup):
                 return "User not authorized to call the tool"
 
             # Check if this is the default user - if so, use base client directly
-            if (not function_group._shared_auth_provider
-                    or session_id == function_group._shared_auth_provider.config.default_user_id):
+            if (not function_group._shared_auth_provider or session_id == function_group._default_user_id):
                 # Use base client directly for default user
                 client = function_group.mcp_client
                 session_tool = await client.get_tool(tool.name)
@@ -507,7 +508,9 @@ async def mcp_client_function_group(config: MCPClientConfig, _builder: Builder):
                               reconnect_max_backoff=config.reconnect_max_backoff)
     elif config.server.transport == "streamable-http":
         # Use default_user_id for the base client
-        base_user_id = auth_provider.config.default_user_id if auth_provider else None
+        # For interactive OAuth2: from config. For service accounts: defaults to server URL
+        base_user_id = getattr(auth_provider.config, 'default_user_id', str(
+            config.server.url)) if auth_provider else None
         client = MCPStreamableHTTPClient(str(config.server.url),
                                          auth_provider=auth_provider,
                                          user_id=base_user_id,
@@ -529,6 +532,18 @@ async def mcp_client_function_group(config: MCPClientConfig, _builder: Builder):
     group._shared_auth_provider = auth_provider
     group._client_config = config
 
+    # Set auth provider config defaults
+    # For interactive OAuth2: use config values
+    # For service accounts: default_user_id = server URL, allow_default_user_id_for_tool_calls = True
+    if auth_provider:
+        group._default_user_id = getattr(auth_provider.config, 'default_user_id', str(config.server.url))
+        group._allow_default_user_id_for_tool_calls = getattr(auth_provider.config,
+                                                              'allow_default_user_id_for_tool_calls',
+                                                              True)
+    else:
+        group._default_user_id = None
+        group._allow_default_user_id_for_tool_calls = True
+
     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.

nat/plugins/mcp/utils.py

@@ -47,7 +47,7 @@ def model_from_mcp_schema(name: str, mcp_input_schema: dict) -> type[BaseModel]:
         "integer": int,
         "boolean": bool,
         "array": list,
-        "null": None,
+        "null": type(None),
         "object": dict,
     }
 
@@ -58,51 +58,168 @@ def model_from_mcp_schema(name: str, mcp_input_schema: dict) -> type[BaseModel]:
     def _generate_valid_classname(class_name: str):
         return class_name.replace('_', ' ').replace('-', ' ').title().replace(' ', '')
 
-    def _generate_field(field_name: str, field_properties: dict[str, Any]) -> tuple:
-        json_type = field_properties.get("type", "string")
-        enum_vals = field_properties.get("enum")
-
+    def _resolve_schema_type(schema: dict[str, Any], name: str) -> Any:
+        """
+        Recursively resolve a JSON schema to a Python type.
+        Handles nested anyOf/oneOf, arrays, objects, enums, and primitive types.
+        """
+        # Check for anyOf/oneOf first
+        any_of = schema.get("anyOf")
+        one_of = schema.get("oneOf")
+
+        if any_of or one_of:
+            union_schemas = any_of if any_of else one_of
+            resolved_type: Any = None
+
+            if union_schemas:
+                for sub_schema in union_schemas:
+                    mapped = _resolve_schema_type(sub_schema, name)
+                    if resolved_type is None:
+                        resolved_type = mapped
+                    elif mapped is not type(None):
+                        # Don't add None here, handle separately
+                        resolved_type = resolved_type | mapped
+                    else:
+                        # If we encounter null, combine with None at the end
+                        resolved_type = resolved_type | None if resolved_type else type(None)
+
+            return resolved_type if resolved_type is not None else Any
+
+        # Handle enum values
+        enum_vals = schema.get("enum")
         if enum_vals:
-            enum_name = f"{field_name.capitalize()}Enum"
-            field_type = Enum(enum_name, {item: item for item in enum_vals})
-
-        elif json_type == "object" and "properties" in field_properties:
-            field_type = model_from_mcp_schema(name=field_name, mcp_input_schema=field_properties)
-        elif json_type == "array" and "items" in field_properties:
-            item_properties = field_properties.get("items", {})
-            if item_properties.get("type") == "object":
-                item_type = model_from_mcp_schema(name=field_name, mcp_input_schema=item_properties)
+            # Check if enum contains null
+            has_null = any(val is None or val == "null" for val in enum_vals)
+            # Filter out None/null values from enum
+            non_null_vals = [v for v in enum_vals if v is not None and v != "null"]
+
+            if non_null_vals:
+                enum_name = f"{name.capitalize()}Enum"
+                enum_type: Any = Enum(enum_name, {item: item for item in non_null_vals})
+                # If enum had null, make it a union with None
+                return enum_type | None if has_null else enum_type
+            elif has_null:
+                # Enum only contains null
+                return type(None)
             else:
-                item_type = _type_map.get(item_properties.get("type", "string"), Any)
-            field_type = list[item_type]
-        elif isinstance(json_type, list):
-            field_type = None
-            for t in json_type:
-                mapped = _type_map.get(t, Any)
-                field_type = mapped if field_type is None else field_type | mapped
-
-            return field_type, Field(
-                default=field_properties.get("default", None if "null" in json_type else ...),
-                description=field_properties.get("description", "")
-            )
-        else:
-            field_type = _type_map.get(json_type, Any)
+                # Empty enum (shouldn't happen but handle gracefully)
+                return Any
+
+        schema_type = schema.get("type")
+
+        # Handle type as list (e.g., ["string", "integer", "null"])
+        if isinstance(schema_type, list):
+            list_type: Any = None
+            for t in schema_type:
+                if t == "array":
+                    # Incorporate the mapped type of items
+                    item_schema = schema.get("items", {})
+                    if item_schema:
+                        item_type = _resolve_schema_type(item_schema, name)
+                        mapped = list[item_type]
+                    else:
+                        mapped = _type_map.get(t, Any)
+                elif t == "object":
+                    # Incorporate the mapped type from properties
+                    if "properties" in schema:
+                        mapped = model_from_mcp_schema(name=name, mcp_input_schema=schema)
+                    else:
+                        mapped = _type_map.get(t, Any)
+                else:
+                    mapped = _type_map.get(t, Any)
+
+                list_type = mapped if list_type is None else list_type | mapped
+            return list_type if list_type is not None else Any
+
+        # Handle null type
+        if schema_type == "null":
+            return type(None)
+
+        # Handle object type
+        if schema_type == "object" and "properties" in schema:
+            return model_from_mcp_schema(name=name, mcp_input_schema=schema)
+
+        # Handle array type
+        if schema_type == "array" and "items" in schema:
+            item_schema = schema.get("items", {})
+            # Recursively resolve item type (handles nested anyOf/oneOf)
+            item_type = _resolve_schema_type(item_schema, name)
+            return list[item_type]
+
+        # Handle primitive types
+        if schema_type is not None:
+            return _type_map.get(schema_type, Any)
+
+        return Any
+
+    def _has_null_in_type(field_properties: dict[str, Any]) -> bool:
+        """Check if a schema contains null as a valid type."""
+        # Check anyOf/oneOf for null
+        any_of = field_properties.get("anyOf")
+        one_of = field_properties.get("oneOf")
+        if any_of or one_of:
+            union_schemas = any_of if any_of else one_of
+            if union_schemas:
+                for schema in union_schemas:
+                    if schema.get("type") == "null":
+                        return True
+
+        # Check type list for null
+        json_type = field_properties.get("type")
+        if isinstance(json_type, list) and "null" in json_type:
+            return True
+
+        # Check enum for null (Python None or string "null")
+        enum_vals = field_properties.get("enum")
+        if enum_vals:
+            for val in enum_vals:
+                if val is None or val == "null":
+                    return True
+
+        # Check const for null (Python None or string "null")
+        if "const" in field_properties:
+            const_val = field_properties.get("const")
+            if const_val is None or const_val == "null":
+                return True
+
+        return False
+
+    def _generate_field(field_name: str, field_properties: dict[str, Any]) -> tuple:
+        """
+        Generate a Pydantic field from JSON schema properties.
+        Uses _resolve_schema_type for type resolution and handles field-specific logic.
+        """
+        # Resolve the field type using the unified resolver
+        field_type = _resolve_schema_type(field_properties, field_name)
+
+        # Check if the type includes null
+        has_null = _has_null_in_type(field_properties)
 
         # Determine the default value based on whether the field is required
+        default_value = field_properties.get("default")
+
         if field_name in required_fields:
-            # Field is required - use explicit default if provided, otherwise make it required
-            default_value = field_properties.get("default", ...)
+            # Field is required - use explicit default if provided, otherwise use ... to enforce presence
+            if default_value is None and "default" not in field_properties:
+                # Required field without explicit default: always use ... even if nullable
+                default_value = ...
+            # Make the field type nullable if it allows null
+            if has_null:
+                field_type = field_type | None
         else:
             # Field is optional - use explicit default if provided, otherwise None
-            default_value = field_properties.get("default", None)
-            # Make the type optional if no default was provided
-            if "default" not in field_properties:
+            if default_value is None:
+                default_value = None
+            # Make the type optional if no default was provided and not already nullable
+            if "default" not in field_properties and not has_null:
                 field_type = field_type | None
 
+        # Handle nullable property (less common, but still supported)
         nullable = field_properties.get("nullable", False)
-        description = field_properties.get("description", "")
+        if nullable and not has_null:
+            field_type = field_type | None
 
-        field_type = field_type | None if nullable else field_type
+        description = field_properties.get("description", "")
 
         return field_type, Field(default=default_value, description=description)
 

{nvidia_nat_mcp-1.3.0rc6.dist-info → nvidia_nat_mcp-1.4.0a20251123.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nvidia-nat-mcp
-Version: 1.3.0rc6
+Version: 1.4.0a20251123
 Summary: Subpackage for MCP client integration in NeMo Agent toolkit
 Author: NVIDIA Corporation
 Maintainer: NVIDIA Corporation
@@ -16,7 +16,7 @@ Requires-Python: <3.14,>=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE-3rd-party.txt
 License-File: LICENSE.md
-Requires-Dist: nvidia-nat==v1.3.0-rc6
+Requires-Dist: nvidia-nat==v1.4.0a20251123
 Requires-Dist: aiorwlock~=1.5
 Requires-Dist: mcp~=1.14
 Dynamic: license-file

nvidia_nat_mcp-1.4.0a20251123.dist-info/RECORD

@@ -0,0 +1,27 @@
+nat/meta/pypi.md,sha256=EYyJTCCEOWzuuz-uNaYJ_WBk55Jiig87wcUr9E4g0yw,1484
+nat/plugins/mcp/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
+nat/plugins/mcp/client_base.py,sha256=Jm7FX4V5125vmqX9unAG4q-eeqgTyzfHgw14sZqnTRM,26753
+nat/plugins/mcp/client_config.py,sha256=l9tVUHe8WdFPJ9rXDg8dZkQi1dvHGYwoqQ8Glqg2LGs,6783
+nat/plugins/mcp/client_impl.py,sha256=BV7r39ijipAz9foAKtrLTairkaOjDH8Bnluw8sf72Ek,27902
+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=xNfBIF__ugJKFEjkYEM417wWM1PpuTaCMGtSFmxHSuA,6089
+nat/plugins/mcp/utils.py,sha256=dUIig7jeKz0ctb4o38jFGbe2uvM3DMR3PSJjfN_Lr5M,9111
+nat/plugins/mcp/auth/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
+nat/plugins/mcp/auth/auth_flow_handler.py,sha256=v21IK3IKZ2TLEP6wO9r-sJQiilWPq7Ry40M96SAxQFA,9125
+nat/plugins/mcp/auth/auth_provider.py,sha256=BgH66DlZgzhLDLO4cBERpHvNAmli5fMo_SCy11W9aBU,21251
+nat/plugins/mcp/auth/auth_provider_config.py,sha256=ZdiUObYU_Oj8KDDZ-JqkS6kJgup5EBy2ZREUOBw_kkI,4143
+nat/plugins/mcp/auth/register.py,sha256=miNZNmNszGgMCOCADLpH0Nz1nugkDXqdVc7ooapBx-c,1754
+nat/plugins/mcp/auth/token_storage.py,sha256=aS13ZvEJXcYzkZ0GSbrSor4i5bpjD5BkXHQw1iywC9k,9240
+nat/plugins/mcp/auth/service_account/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
+nat/plugins/mcp/auth/service_account/provider.py,sha256=jJmt7PA_U3C59K5chSzmEIOd2bRxRKqXly9vHkTmsHQ,5915
+nat/plugins/mcp/auth/service_account/provider_config.py,sha256=KhvQpROporbNKDFnx8np8b31QhdT3s5zYXd_IVTLgfE,5327
+nat/plugins/mcp/auth/service_account/token_client.py,sha256=PRJ3u8Ts1tpuLuQKOyHDPOIAynSKxW9a928Rhh_4fhc,5981
+nvidia_nat_mcp-1.4.0a20251123.dist-info/licenses/LICENSE-3rd-party.txt,sha256=fOk5jMmCX9YoKWyYzTtfgl-SUy477audFC5hNY4oP7Q,284609
+nvidia_nat_mcp-1.4.0a20251123.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+nvidia_nat_mcp-1.4.0a20251123.dist-info/METADATA,sha256=52-egpVHr45VW2p1147j9s4ZSxnY1UO2tXgJDHR5FOI,2319
+nvidia_nat_mcp-1.4.0a20251123.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_mcp-1.4.0a20251123.dist-info/entry_points.txt,sha256=rYvUp4i-klBr3bVNh7zYOPXret704vTjvCk1qd7FooI,97
+nvidia_nat_mcp-1.4.0a20251123.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
+nvidia_nat_mcp-1.4.0a20251123.dist-info/RECORD,,

nvidia_nat_mcp-1.3.0rc6.dist-info/RECORD

@@ -1,23 +0,0 @@
-nat/meta/pypi.md,sha256=EYyJTCCEOWzuuz-uNaYJ_WBk55Jiig87wcUr9E4g0yw,1484
-nat/plugins/mcp/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
-nat/plugins/mcp/client_base.py,sha256=JIyO2ZJsVkQ1g5BOU2zKXGHg_0yxv16g7_YJAqdCXTA,26504
-nat/plugins/mcp/client_config.py,sha256=l9tVUHe8WdFPJ9rXDg8dZkQi1dvHGYwoqQ8Glqg2LGs,6783
-nat/plugins/mcp/client_impl.py,sha256=j7cKAUBKtZAY3mt5Mm8VqgqMhRZk7kzvUd1nwMU_h0o,27072
-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=xNfBIF__ugJKFEjkYEM417wWM1PpuTaCMGtSFmxHSuA,6089
-nat/plugins/mcp/utils.py,sha256=4kNF5FJRiDUn-3fQcsvwvWtG6tYG1y4jU7vpptp0fsA,4522
-nat/plugins/mcp/auth/__init__.py,sha256=GUJrgGtpvyMUCjUBvR3faAdv-tZzbU9W-izgx9aMEQg,680
-nat/plugins/mcp/auth/auth_flow_handler.py,sha256=v21IK3IKZ2TLEP6wO9r-sJQiilWPq7Ry40M96SAxQFA,9125
-nat/plugins/mcp/auth/auth_provider.py,sha256=BgH66DlZgzhLDLO4cBERpHvNAmli5fMo_SCy11W9aBU,21251
-nat/plugins/mcp/auth/auth_provider_config.py,sha256=b1AaXzOuAkygKXAWSxMKWg8wfW8k33tmUUq6Dk5Mmwk,4038
-nat/plugins/mcp/auth/register.py,sha256=L2x69NjJPS4s6CCE5myzWVrWn3e_ttHyojmGXvBipMg,1228
-nat/plugins/mcp/auth/token_storage.py,sha256=aS13ZvEJXcYzkZ0GSbrSor4i5bpjD5BkXHQw1iywC9k,9240
-nvidia_nat_mcp-1.3.0rc6.dist-info/licenses/LICENSE-3rd-party.txt,sha256=fOk5jMmCX9YoKWyYzTtfgl-SUy477audFC5hNY4oP7Q,284609
-nvidia_nat_mcp-1.3.0rc6.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-nvidia_nat_mcp-1.3.0rc6.dist-info/METADATA,sha256=yTrTMrdDdfjXpxpGKrLbSI5xyZ_xBwk8lYrp0BO1c-A,2308
-nvidia_nat_mcp-1.3.0rc6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nvidia_nat_mcp-1.3.0rc6.dist-info/entry_points.txt,sha256=rYvUp4i-klBr3bVNh7zYOPXret704vTjvCk1qd7FooI,97
-nvidia_nat_mcp-1.3.0rc6.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
-nvidia_nat_mcp-1.3.0rc6.dist-info/RECORD,,