nvidia-nat-mcp 1.3.0a20251005__py3-none-any.whl → 1.3.0a20251007__py3-none-any.whl

nat/plugins/mcp/client_config.py

@@ -0,0 +1,131 @@
+# 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.
+
+from datetime import timedelta
+from typing import Literal
+
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic import HttpUrl
+from pydantic import model_validator
+
+from nat.data_models.component_ref import AuthenticationRef
+from nat.data_models.function import FunctionGroupBaseConfig
+
+
+class MCPToolOverrideConfig(BaseModel):
+    """
+    Configuration for overriding tool properties when exposing from MCP server.
+    """
+    alias: str | None = Field(default=None, description="Override the tool name (function name in the workflow)")
+    description: str | None = Field(default=None, description="Override the tool description")
+
+
+class MCPServerConfig(BaseModel):
+    """
+    Server connection details for MCP client.
+    Supports stdio, sse, and streamable-http transports.
+    streamable-http is the recommended default for HTTP-based connections.
+    """
+    transport: Literal["stdio", "sse", "streamable-http"] = Field(
+        ..., description="Transport type to connect to the MCP server (stdio, sse, or streamable-http)")
+    url: HttpUrl | None = Field(default=None,
+                                description="URL of the MCP server (for sse or streamable-http transport)")
+    command: str | None = Field(default=None,
+                                description="Command to run for stdio transport (e.g. 'python' or 'docker')")
+    args: list[str] | None = Field(default=None, description="Arguments for the stdio command")
+    env: dict[str, str] | None = Field(default=None, description="Environment variables for the stdio process")
+
+    # Authentication configuration
+    auth_provider: str | AuthenticationRef | None = Field(default=None,
+                                                          description="Reference to authentication provider")
+
+    @model_validator(mode="after")
+    def validate_model(self):
+        """Validate that stdio and SSE/Streamable HTTP properties are mutually exclusive."""
+        if self.transport == "stdio":
+            if self.url is not None:
+                raise ValueError("url should not be set when using stdio transport")
+            if not self.command:
+                raise ValueError("command is required when using stdio transport")
+            # Auth is not supported for stdio transport
+            if self.auth_provider is not None:
+                raise ValueError("Authentication is not supported for stdio transport")
+        elif self.transport == "sse":
+            if self.command is not None or self.args is not None or self.env is not None:
+                raise ValueError("command, args, and env should not be set when using sse transport")
+            if not self.url:
+                raise ValueError("url is required when using sse transport")
+            # Auth is not supported for SSE transport
+            if self.auth_provider is not None:
+                raise ValueError("Authentication is not supported for SSE transport.")
+        elif self.transport == "streamable-http":
+            if self.command is not None or self.args is not None or self.env is not None:
+                raise ValueError("command, args, and env should not be set when using streamable-http transport")
+            if not self.url:
+                raise ValueError("url is required when using streamable-http transport")
+
+        return self
+
+
+class MCPClientConfig(FunctionGroupBaseConfig, name="mcp_client"):
+    """
+    Configuration for connecting to an MCP server as a client and exposing selected tools.
+    """
+    server: MCPServerConfig = Field(..., description="Server connection details (transport, url/command, etc.)")
+    tool_call_timeout: timedelta = Field(
+        default=timedelta(seconds=60),
+        description="Timeout (in seconds) for the MCP tool call. Defaults to 60 seconds.")
+    auth_flow_timeout: timedelta = Field(
+        default=timedelta(seconds=300),
+        description="Timeout (in seconds) for the MCP auth flow. When the tool call requires interactive \
+        authentication, this timeout is used. Defaults to 300 seconds.")
+    reconnect_enabled: bool = Field(
+        default=True,
+        description="Whether to enable reconnecting to the MCP server if the connection is lost. \
+        Defaults to True.")
+    reconnect_max_attempts: int = Field(default=2,
+                                        ge=0,
+                                        description="Maximum number of reconnect attempts. Defaults to 2.")
+    reconnect_initial_backoff: float = Field(
+        default=0.5, ge=0.0, description="Initial backoff time for reconnect attempts. Defaults to 0.5 seconds.")
+    reconnect_max_backoff: float = Field(
+        default=50.0, ge=0.0, description="Maximum backoff time for reconnect attempts. Defaults to 50 seconds.")
+    tool_overrides: dict[str, MCPToolOverrideConfig] | None = Field(
+        default=None,
+        description="""Optional tool name overrides and description changes.
+        Example:
+          tool_overrides:
+            calculator_add:
+              alias: "add_numbers"
+              description: "Add two numbers together"
+            calculator_multiply:
+              description: "Multiply two numbers"  # alias defaults to original name
+        """)
+    session_aware_tools: bool = Field(default=True,
+                                      description="Session-aware tools are created if True. Defaults to True.")
+    max_sessions: int = Field(default=100,
+                              ge=1,
+                              description="Maximum number of concurrent session clients. Defaults to 100.")
+    session_idle_timeout: timedelta = Field(
+        default=timedelta(hours=1),
+        description="Time after which inactive sessions are cleaned up. Defaults to 1 hour.")
+
+    @model_validator(mode="after")
+    def _validate_reconnect_backoff(self) -> "MCPClientConfig":
+        """Validate reconnect backoff values."""
+        if self.reconnect_max_backoff < self.reconnect_initial_backoff:
+            raise ValueError("reconnect_max_backoff must be greater than or equal to reconnect_initial_backoff")
+        return self

nat/plugins/mcp/client_impl.py

@@ -13,29 +13,41 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import logging
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from dataclasses import field
+from datetime import datetime
 from datetime import timedelta
-from typing import Literal
 
+import aiorwlock
 from pydantic import BaseModel
-from pydantic import Field
-from pydantic import HttpUrl
-from pydantic import model_validator
 
+from nat.authentication.interfaces import AuthProviderBase
 from nat.builder.builder import Builder
 from nat.builder.function import FunctionGroup
 from nat.cli.register_workflow import register_function_group
-from nat.data_models.component_ref import AuthenticationRef
-from nat.data_models.function import FunctionGroupBaseConfig
-from nat.plugins.mcp.tool import mcp_tool_function
+from nat.plugins.mcp.client_base import MCPBaseClient
+from nat.plugins.mcp.client_config import MCPClientConfig
+from nat.plugins.mcp.client_config import MCPToolOverrideConfig
+from nat.plugins.mcp.utils import truncate_session_id
 
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class SessionData:
+    """Container for all session-related data."""
+    client: MCPBaseClient
+    last_activity: datetime
+    ref_count: int = 0
+    lock: asyncio.Lock = field(default_factory=asyncio.Lock)
+
+
 class MCPFunctionGroup(FunctionGroup):
     """
-    A specialized FunctionGroup for MCP clients that includes MCP-specific attributes
-    with proper type safety.
+    A specialized FunctionGroup for MCP clients that includes MCP-specific attributes with session management.
     """
 
     def __init__(self, *args, **kwargs):
@@ -45,6 +57,20 @@ class MCPFunctionGroup(FunctionGroup):
         self._mcp_client_server_name: str | None = None
         self._mcp_client_transport: str | None = None
 
+        # Session management - consolidated data structure
+        self._sessions: dict[str, SessionData] = {}
+
+        # Use RWLock for better concurrency: multiple readers (tool calls) can access
+        # existing sessions simultaneously, while writers (create/delete) get exclusive access
+        self._session_rwlock = aiorwlock.RWLock()
+        # Throttled cleanup control
+        self._last_cleanup_check: datetime = datetime.now()
+        self._cleanup_check_interval: timedelta = timedelta(minutes=5)
+
+        # Shared components for session client creation
+        self._shared_auth_provider: AuthProviderBase | None = None
+        self._client_config: MCPClientConfig | None = None
+
     @property
     def mcp_client(self):
         """Get the MCP client instance."""
@@ -75,109 +101,253 @@ class MCPFunctionGroup(FunctionGroup):
         """Set the MCP client transport type."""
         self._mcp_client_transport = transport
 
+    @property
+    def session_count(self) -> int:
+        """Current number of active sessions."""
+        return len(self._sessions)
 
-class MCPToolOverrideConfig(BaseModel):
-    """
-    Configuration for overriding tool properties when exposing from MCP server.
-    """
-    alias: str | None = Field(default=None, description="Override the tool name (function name in the workflow)")
-    description: str | None = Field(default=None, description="Override the tool description")
-
-
-class MCPServerConfig(BaseModel):
-    """
-    Server connection details for MCP client.
-    Supports stdio, sse, and streamable-http transports.
-    streamable-http is the recommended default for HTTP-based connections.
-    """
-    transport: Literal["stdio", "sse", "streamable-http"] = Field(
-        ..., description="Transport type to connect to the MCP server (stdio, sse, or streamable-http)")
-    url: HttpUrl | None = Field(default=None,
-                                description="URL of the MCP server (for sse or streamable-http transport)")
-    command: str | None = Field(default=None,
-                                description="Command to run for stdio transport (e.g. 'python' or 'docker')")
-    args: list[str] | None = Field(default=None, description="Arguments for the stdio command")
-    env: dict[str, str] | None = Field(default=None, description="Environment variables for the stdio process")
-
-    # Authentication configuration
-    auth_provider: str | AuthenticationRef | None = Field(default=None,
-                                                          description="Reference to authentication provider")
-
-    @model_validator(mode="after")
-    def validate_model(self):
-        """Validate that stdio and SSE/Streamable HTTP properties are mutually exclusive."""
-        if self.transport == "stdio":
-            if self.url is not None:
-                raise ValueError("url should not be set when using stdio transport")
-            if not self.command:
-                raise ValueError("command is required when using stdio transport")
-            # Auth is not supported for stdio transport
-            if self.auth_provider is not None:
-                raise ValueError("Authentication is not supported for stdio transport")
-        elif self.transport == "sse":
-            if self.command is not None or self.args is not None or self.env is not None:
-                raise ValueError("command, args, and env should not be set when using sse transport")
-            if not self.url:
-                raise ValueError("url is required when using sse transport")
-            # Auth is not supported for SSE transport
-            if self.auth_provider is not None:
-                raise ValueError("Authentication is not supported for SSE transport.")
-        elif self.transport == "streamable-http":
-            if self.command is not None or self.args is not None or self.env is not None:
-                raise ValueError("command, args, and env should not be set when using streamable-http transport")
-            if not self.url:
-                raise ValueError("url is required when using streamable-http transport")
-
-        return self
-
-
-class MCPClientConfig(FunctionGroupBaseConfig, name="mcp_client"):
-    """
-    Configuration for connecting to an MCP server as a client and exposing selected tools.
+    @property
+    def session_limit(self) -> int:
+        """Maximum allowed sessions."""
+        return self._client_config.max_sessions if self._client_config else 100
+
+    def _get_session_id_from_context(self) -> str | None:
+        """Get the session ID from the current context."""
+        try:
+            from nat.builder.context import Context as _Ctx
+
+            # Get session id from context, authentication is done per-websocket session for tool calls
+            session_id = None
+            cookies = getattr(_Ctx.get().metadata, "cookies", None)
+            if cookies:
+                session_id = cookies.get("nat-session")
+
+            if not session_id:
+                # use default user id if allowed
+                if self._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
+            return session_id
+        except Exception:
+            return None
+
+    async def cleanup_sessions(self, max_age: timedelta | None = None) -> int:
+        """
+        Manually trigger cleanup of inactive sessions.
+
+        Args:
+            max_age: Maximum age for sessions before cleanup. If None, uses configured timeout.
+
+        Returns:
+            Number of sessions cleaned up.
+        """
+        sessions_before = len(self._sessions)
+        await self._cleanup_inactive_sessions(max_age)
+        sessions_after = len(self._sessions)
+        return sessions_before - sessions_after
+
+    async def _cleanup_inactive_sessions(self, max_age: timedelta | None = None):
+        """Remove clients for sessions inactive longer than max_age.
+
+        This method uses the RWLock writer to ensure thread-safe cleanup.
+        """
+        if max_age is None:
+            max_age = self._client_config.session_idle_timeout if self._client_config else timedelta(hours=1)
+
+        async with self._session_rwlock.writer:
+            current_time = datetime.now()
+            inactive_sessions = []
+
+            for session_id, session_data in self._sessions.items():
+                # Skip cleanup if session is actively being used
+                if session_data.ref_count > 0:
+                    continue
+
+                if current_time - session_data.last_activity > max_age:
+                    inactive_sessions.append(session_id)
+
+            for session_id in inactive_sessions:
+                try:
+                    logger.info("Cleaning up inactive session client: %s", truncate_session_id(session_id))
+                    session_data = self._sessions[session_id]
+                    # Close the client connection
+                    await session_data.client.__aexit__(None, None, None)
+                    logger.info("Cleaned up inactive session client: %s", truncate_session_id(session_id))
+                except Exception as e:
+                    logger.warning("Error cleaning up session client %s: %s", truncate_session_id(session_id), e)
+                finally:
+                    # Always remove from tracking to prevent leaks, even if close failed
+                    self._sessions.pop(session_id, None)
+                    logger.info("Cleaned up session tracking for: %s", truncate_session_id(session_id))
+                    logger.info(" Total sessions: %d", len(self._sessions))
+
+    async def _get_session_client(self, session_id: str) -> MCPBaseClient:
+        """Get the appropriate MCP client for the session."""
+        # Throttled cleanup on access
+        now = datetime.now()
+        if now - self._last_cleanup_check > self._cleanup_check_interval:
+            await self._cleanup_inactive_sessions()
+            self._last_cleanup_check = now
+
+        # 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:
+                return self.mcp_client
+
+        # Fast path: check if session already exists (reader lock for concurrent access)
+        async with self._session_rwlock.reader:
+            if session_id in self._sessions:
+                # Update last activity for existing client
+                self._sessions[session_id].last_activity = datetime.now()
+                return self._sessions[session_id].client
+
+        # Check session limit before creating new client (outside writer lock to avoid deadlock)
+        if self._client_config and len(self._sessions) >= self._client_config.max_sessions:
+            # Try cleanup first to free up space
+            await self._cleanup_inactive_sessions()
+
+        # Slow path: create session with writer lock for exclusive access
+        async with self._session_rwlock.writer:
+            # Double-check after acquiring writer lock (another coroutine might have created it)
+            if session_id in self._sessions:
+                self._sessions[session_id].last_activity = datetime.now()
+                return self._sessions[session_id].client
+
+            # Re-check session limit inside writer lock
+            if self._client_config and len(self._sessions) >= self._client_config.max_sessions:
+                logger.warning("Session limit reached (%d), rejecting new session: %s",
+                               self._client_config.max_sessions,
+                               truncate_session_id(session_id))
+                raise RuntimeError(f"Service temporarily unavailable: Maximum concurrent sessions "
+                                   f"({self._client_config.max_sessions}) exceeded. Please try again later.")
+
+            # Create session client lazily
+            logger.info("Creating new MCP client for session: %s", truncate_session_id(session_id))
+            session_client = await self._create_session_client(session_id)
+
+            # Create session data with all components
+            session_data = SessionData(client=session_client, last_activity=datetime.now(), ref_count=0)
+
+            # Cache the session data
+            self._sessions[session_id] = session_data
+            logger.info(" Total sessions: %d", len(self._sessions))
+            return session_client
+
+    @asynccontextmanager
+    async def _session_usage_context(self, session_id: str):
+        """Context manager to track active session usage and prevent cleanup."""
+        # Ensure session exists - create it if it doesn't
+        if session_id not in self._sessions:
+            # Create session client first
+            await self._get_session_client(session_id)
+            # Session should now exist in _sessions
+
+        # Get session data (session must exist at this point)
+        session_data = self._sessions[session_id]
+
+        # Thread-safe reference counting using per-session lock
+        async with session_data.lock:
+            session_data.ref_count += 1
+
+        try:
+            yield
+        finally:
+            async with session_data.lock:
+                session_data.ref_count -= 1
+
+    async def _create_session_client(self, session_id: str) -> MCPBaseClient:
+        """Create a new MCP client instance for the session."""
+        from nat.plugins.mcp.client_base import MCPStreamableHTTPClient
+
+        config = self._client_config
+        if not config:
+            raise RuntimeError("Client config not initialized")
+
+        if config.server.transport == "streamable-http":
+            client = MCPStreamableHTTPClient(
+                str(config.server.url),
+                auth_provider=self._shared_auth_provider,
+                user_id=session_id,  # Pass session_id as user_id for cache isolation
+                tool_call_timeout=config.tool_call_timeout,
+                auth_flow_timeout=config.auth_flow_timeout,
+                reconnect_enabled=config.reconnect_enabled,
+                reconnect_max_attempts=config.reconnect_max_attempts,
+                reconnect_initial_backoff=config.reconnect_initial_backoff,
+                reconnect_max_backoff=config.reconnect_max_backoff)
+        else:
+            # per-user sessions are only supported for streamable-http transport
+            raise ValueError(f"Unsupported transport: {config.server.transport}")
+
+        # Initialize the client
+        await client.__aenter__()
+
+        logger.info("Created session client for session: %s", truncate_session_id(session_id))
+        return client
+
+
+def mcp_session_tool_function(tool, function_group: MCPFunctionGroup):
+    """Create a session-aware NAT function for an MCP tool.
+
+    Routes each invocation to the appropriate per-session MCP client while
+    preserving the original tool input schema, converters, and description.
     """
-    server: MCPServerConfig = Field(..., description="Server connection details (transport, url/command, etc.)")
-    tool_call_timeout: timedelta = Field(
-        default=timedelta(seconds=60),
-        description="Timeout (in seconds) for the MCP tool call. Defaults to 60 seconds.")
-    auth_flow_timeout: timedelta = Field(
-        default=timedelta(seconds=300),
-        description="Timeout (in seconds) for the MCP auth flow. When the tool call requires interactive \
-        authentication, this timeout is used. Defaults to 300 seconds.")
-    reconnect_enabled: bool = Field(
-        default=True,
-        description="Whether to enable reconnecting to the MCP server if the connection is lost. \
-        Defaults to True.")
-    reconnect_max_attempts: int = Field(default=2,
-                                        ge=0,
-                                        description="Maximum number of reconnect attempts. Defaults to 2.")
-    reconnect_initial_backoff: float = Field(
-        default=0.5, ge=0.0, description="Initial backoff time for reconnect attempts. Defaults to 0.5 seconds.")
-    reconnect_max_backoff: float = Field(
-        default=50.0, ge=0.0, description="Maximum backoff time for reconnect attempts. Defaults to 50 seconds.")
-    tool_overrides: dict[str, MCPToolOverrideConfig] | None = Field(
-        default=None,
-        description="""Optional tool name overrides and description changes.
-        Example:
-          tool_overrides:
-            calculator_add:
-              alias: "add_numbers"
-              description: "Add two numbers together"
-            calculator_multiply:
-              description: "Multiply two numbers"  # alias defaults to original name
-        """)
-
-    @model_validator(mode="after")
-    def _validate_reconnect_backoff(self) -> "MCPClientConfig":
-        """Validate reconnect backoff values."""
-        if self.reconnect_max_backoff < self.reconnect_initial_backoff:
-            raise ValueError("reconnect_max_backoff must be greater than or equal to reconnect_initial_backoff")
-        return self
+    from nat.builder.function import FunctionInfo
+
+    def _convert_from_str(input_str: str) -> tool.input_schema:
+        return tool.input_schema.model_validate_json(input_str)
+
+    async def _response_fn(tool_input: BaseModel | None = None, **kwargs) -> str:
+        """Response function for the session-aware tool."""
+        try:
+            # Route to the appropriate session client
+            session_id = function_group._get_session_id_from_context()
+
+            # If no session is available and default-user fallback is disabled, deny the call
+            if function_group._shared_auth_provider and session_id is None:
+                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):
+                # Use base client directly for default user
+                client = function_group.mcp_client
+                session_tool = await client.get_tool(tool.name)
+            else:
+                # Use session usage context to prevent cleanup during tool execution
+                async with function_group._session_usage_context(session_id):
+                    client = await function_group._get_session_client(session_id)
+                    session_tool = await client.get_tool(tool.name)
+
+            # Preserve original calling convention
+            if tool_input:
+                args = tool_input.model_dump()
+                return await session_tool.acall(args)
+
+            _ = session_tool.input_schema.model_validate(kwargs)
+            return await session_tool.acall(kwargs)
+        except Exception as e:
+            if tool_input:
+                logger.warning("Error calling tool %s with serialized input: %s",
+                               tool.name,
+                               tool_input.model_dump(),
+                               exc_info=True)
+            else:
+                logger.warning("Error calling tool %s with input: %s", tool.name, kwargs, exc_info=True)
+            return str(e)
+
+    return FunctionInfo.create(single_fn=_response_fn,
+                               description=tool.description,
+                               input_schema=tool.input_schema,
+                               converters=[_convert_from_str])
 
 
 @register_function_group(config_type=MCPClientConfig)
 async def mcp_client_function_group(config: MCPClientConfig, _builder: Builder):
     """
     Connect to an MCP server and expose tools as a function group.
+
     Args:
         config: The configuration for the MCP client
         _builder: The builder
@@ -215,8 +385,11 @@ async def mcp_client_function_group(config: MCPClientConfig, _builder: Builder):
                               reconnect_initial_backoff=config.reconnect_initial_backoff,
                               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
         client = MCPStreamableHTTPClient(str(config.server.url),
                                          auth_provider=auth_provider,
+                                         user_id=base_user_id,
                                          tool_call_timeout=config.tool_call_timeout,
                                          auth_flow_timeout=config.auth_flow_timeout,
                                          reconnect_enabled=config.reconnect_enabled,
@@ -231,6 +404,10 @@ async def mcp_client_function_group(config: MCPClientConfig, _builder: Builder):
     # Create the MCP function group
     group = MCPFunctionGroup(config=config)
 
+    # Store shared components for session client creation
+    group._shared_auth_provider = auth_provider
+    group._client_config = config
+
     async with client:
         # Expose the live MCP client on the function group instance so other components (e.g., HTTP endpoints)
         # can reuse the already-established session instead of creating a new client per request.
@@ -250,13 +427,17 @@ async def mcp_client_function_group(config: MCPClientConfig, _builder: Builder):
             function_name = override.alias if override and override.alias else tool_name
             description = override.description if override and override.description else tool.description
 
-            # Create the tool function
-            tool_fn = mcp_tool_function(tool)
+            # Create the tool function according to configuration
+            if config.session_aware_tools:
+                tool_fn = mcp_session_tool_function(tool, group)
+            else:
+                from nat.plugins.mcp.tool import mcp_tool_function
+                tool_fn = mcp_tool_function(tool)
 
             # Normalize optional typing for linter/type-checker compatibility
             single_fn = tool_fn.single_fn
             if single_fn is None:
-                # Should not happen because mcp_tool_function always sets a single_fn
+                # Should not happen because FunctionInfo always sets a single_fn
                 logger.warning("Skipping tool %s because single_fn is None", function_name)
                 continue
 
@@ -280,6 +461,7 @@ def mcp_apply_tool_alias_and_description(
         all_tools: dict, tool_overrides: dict[str, MCPToolOverrideConfig] | None) -> dict[str, MCPToolOverrideConfig]:
     """
     Filter tool overrides to only include tools that exist in the MCP server.
+
     Args:
         all_tools: The tools from the MCP server
         tool_overrides: The tool overrides to apply

nat/plugins/mcp/tool.py

@@ -26,6 +26,7 @@ from nat.builder.function_info import FunctionInfo
 from nat.cli.register_workflow import register_function
 from nat.data_models.function import FunctionBaseConfig
 from nat.plugins.mcp.client_base import MCPToolClient
+from nat.utils.decorators import deprecated
 
 logger = logging.getLogger(__name__)
 
@@ -109,6 +110,10 @@ def mcp_tool_function(tool: MCPToolClient) -> FunctionInfo:
 
 
 @register_function(config_type=MCPToolConfig)
+@deprecated(
+    reason=
+    "This function is being replaced with the new mcp_client function group that supports additional MCP features",
+    feature_name="mcp_tool_wrapper")
 async def mcp_tool(config: MCPToolConfig, builder: Builder):
     """
     Generate a NeMo Agent Toolkit Function that wraps a tool provided by the MCP server.

nat/plugins/mcp/utils.py

@@ -21,6 +21,22 @@ from pydantic import Field
 from pydantic import create_model
 
 
+def truncate_session_id(session_id: str, max_length: int = 10) -> str:
+    """
+    Truncate a session ID for logging purposes.
+
+    Args:
+        session_id: The session ID to truncate
+        max_length: Maximum length before truncation (default: 10)
+
+    Returns:
+        Truncated session ID with "..." if longer than max_length, otherwise full ID
+    """
+    if len(session_id) > max_length:
+        return session_id[:max_length] + "..."
+    return session_id
+
+
 def model_from_mcp_schema(name: str, mcp_input_schema: dict) -> type[BaseModel]:
     """
     Create a pydantic model from the input schema of the MCP tool