nvidia-nat-mcp 1.4.0a20260107__py3-none-any.whl

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

@@ -0,0 +1,136 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, 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-2026, 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-2026, 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/auth/token_storage.py

@@ -0,0 +1,265 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, 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 the built-in object store provided by the NeMo Agent toolkit.
+
+    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()