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

nat/plugins/mcp/auth/token_storage.py

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

nat/plugins/mcp/client_base.py

@@ -16,15 +16,16 @@
 from __future__ import annotations
 
 import asyncio
-import json
 import logging
 from abc import ABC
 from abc import abstractmethod
 from collections.abc import AsyncGenerator
+from collections.abc import Callable
 from contextlib import AsyncExitStack
 from contextlib import asynccontextmanager
 from datetime import timedelta
 
+import anyio
 import httpx
 
 from mcp import ClientSession
@@ -33,9 +34,9 @@ from mcp.client.stdio import StdioServerParameters
 from mcp.client.stdio import stdio_client
 from mcp.client.streamable_http import streamablehttp_client
 from mcp.types import TextContent
+from nat.authentication.interfaces import AuthenticatedContext
+from nat.authentication.interfaces import AuthFlowType
 from nat.authentication.interfaces import AuthProviderBase
-from nat.data_models.authentication import AuthReason
-from nat.data_models.authentication import AuthRequest
 from nat.plugins.mcp.exception_handler import convert_to_mcp_error
 from nat.plugins.mcp.exception_handler import format_mcp_error
 from nat.plugins.mcp.exception_handler import mcp_exception_handler
@@ -53,74 +54,71 @@ class AuthAdapter(httpx.Auth):
     Converts AuthProviderBase to httpx.Auth interface for dynamic token management.
     """
 
-    def __init__(self, auth_provider: AuthProviderBase, auth_for_tool_calls_only: bool = False):
+    def __init__(self, auth_provider: AuthProviderBase, user_id: str | None = None):
         self.auth_provider = auth_provider
-        self.auth_for_tool_calls_only = auth_for_tool_calls_only
+        self.user_id = user_id  # Session-specific user ID for cache isolation
+        # each adapter instance has its own lock to avoid unnecessary delays for multiple clients
+        self._lock = anyio.Lock()
+        # Track whether we're currently in an interactive authentication flow
+        self.is_authenticating = False
 
     async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]:
         """Add authentication headers to the request using NAT auth provider."""
-        # Check if we should only auth tool calls, Is this needed?
-        if self.auth_for_tool_calls_only and not self._is_tool_call_request(request):
-            # Skip auth for non-tool calls
-            yield request
-            return
-
-        try:
-            # Get fresh auth headers from the NAT auth provider
-            auth_headers = await self._get_auth_headers(reason=AuthReason.NORMAL)
-            request.headers.update(auth_headers)
-        except Exception as e:
-            logger.info("Failed to get auth headers: %s", e)
-            # Continue without auth headers if auth fails
-
-        response = yield request
-
-        # Handle 401 responses by retrying with fresh auth
-        if response.status_code == 401:
+        async with self._lock:
             try:
-                # Get fresh auth headers with 401 context
-                auth_headers = await self._get_auth_headers(reason=AuthReason.RETRY_AFTER_401, response=response)
+                # Get auth headers from the NAT auth provider:
+                # 1. If discovery is yet to done this will return None and request will be sent without auth header.
+                # 2. If discovery is done, this will return the auth header from cache if the token is still valid
+                auth_headers = await self._get_auth_headers(request=request, response=None)
                 request.headers.update(auth_headers)
-                yield request  # Retry the request
             except Exception as e:
-                logger.info("Failed to refresh auth after 401: %s", e)
+                logger.info("Failed to get auth headers: %s", e)
+                # Continue without auth headers if auth fails
+
+            response = yield request
+
+            # Handle 401 responses by retrying with fresh auth
+            if response.status_code == 401:
+                try:
+                    # 401 can happen if:
+                    # 1. The request was sent without auth header
+                    # 2. The auth headers are invalid
+                    # 3. The auth headers are expired
+                    # 4. The auth headers are revoked
+                    # 5. Auth config on the MCP server has changed
+                    # In this case we attempt to re-run discovery and authentication
+
+                    # Signal that we're entering interactive auth flow
+                    self.is_authenticating = True
+                    logger.debug("Starting authentication flow due to 401 response")
+
+                    auth_headers = await self._get_auth_headers(request=request, response=response)
+                    request.headers.update(auth_headers)
+                    yield request  # Retry the request
+                except Exception as e:
+                    logger.info("Failed to refresh auth after 401: %s", e)
+                    raise
+                finally:
+                    # Signal that auth flow is complete
+                    self.is_authenticating = False
+                    logger.debug("Authentication flow completed")
         return
 
-    def _is_tool_call_request(self, request: httpx.Request) -> bool:
-        """Check if this is a tool call request based on the request body."""
-        try:
-            # Check if the request body contains a tool call
-            if request.content:
-                body = json.loads(request.content.decode('utf-8'))
-                # Check if it's a JSON-RPC request with method "tools/call"
-                if (isinstance(body, dict) and body.get("method") == "tools/call"):
-                    return True
-        except (json.JSONDecodeError, UnicodeDecodeError, AttributeError):
-            # If we can't parse the body, assume it's not a tool call
-            pass
-        return False
-
-    async def _get_auth_headers(self, reason: AuthReason, response: httpx.Response | None = None) -> dict[str, str]:
+    async def _get_auth_headers(self,
+                                request: httpx.Request | None = None,
+                                response: httpx.Response | None = None) -> dict[str, str]:
         """Get authentication headers from the NAT auth provider."""
-        # Build auth request
-        www_authenticate = response.headers.get("WWW-Authenticate", None) if response else None
-        auth_request = AuthRequest(
-            reason=reason,
-            www_authenticate=www_authenticate,
-        )
         try:
-            # Mutating the config is not thread-safe, so we need to lock here
-            # Is mutating the config the only way to pass the auth request to the auth provider? This needs
-            # to be re-visited.
-            self.auth_provider.config.auth_request = auth_request
-            auth_result = await self.auth_provider.authenticate()
+            # Use the user_id passed to this AuthAdapter instance
+            auth_result = await self.auth_provider.authenticate(user_id=self.user_id, response=response)
+
             # Check if we have BearerTokenCred
             from nat.data_models.authentication import BearerTokenCred
             if auth_result.credentials and isinstance(auth_result.credentials[0], BearerTokenCred):
                 token = auth_result.credentials[0].token.get_secret_value()
                 return {"Authorization": f"Bearer {token}"}
             else:
-                logger.warning("Auth provider did not return BearerTokenCred")
+                logger.info("Auth provider did not return BearerTokenCred")
                 return {}
         except Exception as e:
             logger.warning("Failed to get auth token: %s", e)
@@ -134,12 +132,20 @@ class MCPBaseClient(ABC):
     Args:
         transport (str): The type of client to use ('sse', 'stdio', or 'streamable-http')
         auth_provider (AuthProviderBase | None): Optional authentication provider for Bearer token injection
+        tool_call_timeout (timedelta): Timeout for tool calls when authentication is not required
+        auth_flow_timeout (timedelta): Extended timeout for tool calls that may require interactive authentication
+        reconnect_enabled (bool): Whether to automatically reconnect on connection failures
+        reconnect_max_attempts (int): Maximum number of reconnection attempts
+        reconnect_initial_backoff (float): Initial backoff delay in seconds for reconnection attempts
+        reconnect_max_backoff (float): Maximum backoff delay in seconds for reconnection attempts
     """
 
     def __init__(self,
                  transport: str = 'streamable-http',
                  auth_provider: AuthProviderBase | None = None,
-                 tool_call_timeout: timedelta = timedelta(seconds=5),
+                 user_id: str | None = None,
+                 tool_call_timeout: timedelta = timedelta(seconds=60),
+                 auth_flow_timeout: timedelta = timedelta(seconds=300),
                  reconnect_enabled: bool = True,
                  reconnect_max_attempts: int = 2,
                  reconnect_initial_backoff: float = 0.5,
@@ -155,9 +161,13 @@ class MCPBaseClient(ABC):
         self._initial_connection = False
 
         # Convert auth provider to AuthAdapter
-        self._httpx_auth = AuthAdapter(auth_provider) if auth_provider else None
+        self._auth_provider = auth_provider
+        # Use provided user_id or fall back to auth provider's default_user_id
+        effective_user_id = user_id or (auth_provider.config.default_user_id if auth_provider else None)
+        self._httpx_auth = AuthAdapter(auth_provider, effective_user_id) if auth_provider else None
 
         self._tool_call_timeout = tool_call_timeout
+        self._auth_flow_timeout = auth_flow_timeout
 
         # Reconnect configuration
         self._reconnect_enabled = reconnect_enabled
@@ -166,6 +176,10 @@ class MCPBaseClient(ABC):
         self._reconnect_max_backoff = reconnect_max_backoff
         self._reconnect_lock: asyncio.Lock = asyncio.Lock()
 
+    @property
+    def auth_provider(self) -> AuthProviderBase | None:
+        return self._auth_provider
+
     @property
     def transport(self) -> str:
         return self._transport
@@ -248,12 +262,25 @@ class MCPBaseClient(ABC):
     async def _with_reconnect(self, coro):
         """
         Execute an awaited operation, reconnecting once on errors.
+        Does not reconnect if the error occurs during an active authentication flow.
         """
         try:
             return await coro()
         except Exception as e:
+            # Check if error happened during active authentication flow
+            if self._httpx_auth and self._httpx_auth.is_authenticating:
+                # Provide specific error message for authentication timeouts
+                if isinstance(e, TimeoutError):
+                    logger.error("Timeout during user authentication flow - user may have abandoned authentication")
+                    raise RuntimeError(
+                        "Authentication timed out. User did not complete authentication in browser within "
+                        f"{self._auth_flow_timeout.total_seconds()} seconds.") from e
+                else:
+                    logger.error("Error during authentication flow: %s", e)
+                    raise
+
+            # Normal error - attempt reconnect if enabled
             if self._reconnect_enabled:
-                logger.warning("MCP Client operation failed. Attempting reconnect: %s", e)
                 try:
                     await self._reconnect()
                 except Exception as reconnect_err:
@@ -262,7 +289,49 @@ class MCPBaseClient(ABC):
                 return await coro()
             raise
 
-    async def get_tools(self):
+    async def _has_cached_auth_token(self) -> bool:
+        """
+        Check if we have a cached, non-expired authentication token.
+
+        Returns:
+            bool: True if we have a valid cached token, False if authentication may be needed
+        """
+        if not self._auth_provider:
+            return True  # No auth needed
+
+        try:
+            # Check if OAuth2 provider has tokens cached
+            if hasattr(self._auth_provider, '_auth_code_provider'):
+                provider = self._auth_provider._auth_code_provider
+                if provider and hasattr(provider, '_authenticated_tokens'):
+                    # Check if we have at least one non-expired token
+                    for auth_result in provider._authenticated_tokens.values():
+                        if not auth_result.is_expired():
+                            return True
+
+            return False
+        except Exception:
+            # If we can't check, assume we need auth to be safe
+            return False
+
+    async def _get_tool_call_timeout(self) -> timedelta:
+        """
+        Determine the appropriate timeout for a tool call based on authentication state.
+
+        Returns:
+            timedelta: auth_flow_timeout if authentication may be needed, tool_call_timeout otherwise
+        """
+        if self._auth_provider:
+            has_token = await self._has_cached_auth_token()
+            timeout = self._tool_call_timeout if has_token else self._auth_flow_timeout
+            if not has_token:
+                logger.debug("Using extended timeout (%s) for potential interactive authentication", timeout)
+            return timeout
+        else:
+            return self._tool_call_timeout
+
+    @mcp_exception_handler
+    async def get_tools(self) -> dict[str, MCPToolClient]:
         """
         Retrieve a dictionary of all tools served by the MCP server.
         Uses unauthenticated session for discovery.
@@ -270,7 +339,16 @@ class MCPBaseClient(ABC):
 
         async def _get_tools():
             session = self._session
-            return await session.list_tools()
+            try:
+                # Add timeout to the list_tools call.
+                # This is needed because MCP SDK does not support timeout for list_tools()
+                with anyio.fail_after(self._tool_call_timeout.total_seconds()):
+                    tools = await session.list_tools()
+            except TimeoutError as e:
+                from nat.plugins.mcp.exceptions import MCPTimeoutError
+                raise MCPTimeoutError(self.server_name, e)
+
+            return tools
 
         try:
             response = await self._with_reconnect(_get_tools)
@@ -284,8 +362,7 @@ class MCPBaseClient(ABC):
                               tool_name=tool.name,
                               tool_description=tool.description,
                               tool_input_schema=tool.inputSchema,
-                              parent_client=self,
-                              tool_call_timeout=self._tool_call_timeout)
+                              parent_client=self)
             for tool in response.tools
         }
 
@@ -314,12 +391,18 @@ class MCPBaseClient(ABC):
             raise MCPToolNotFoundError(tool_name, self.server_name)
         return tool
 
+    def set_user_auth_callback(self, auth_callback: Callable[[AuthFlowType], AuthenticatedContext]):
+        """Set the user authentication callback."""
+        if self._auth_provider and hasattr(self._auth_provider, "_set_custom_auth_callback"):
+            self._auth_provider._set_custom_auth_callback(auth_callback)
+
     @mcp_exception_handler
     async def call_tool(self, tool_name: str, tool_args: dict | None):
 
         async def _call_tool():
             session = self._session
-            return await session.call_tool(tool_name, tool_args, read_timeout_seconds=self._tool_call_timeout)
+            timeout = await self._get_tool_call_timeout()
+            return await session.call_tool(tool_name, tool_args, read_timeout_seconds=timeout)
 
         return await self._with_reconnect(_call_tool)
 
@@ -334,13 +417,15 @@ class MCPSSEClient(MCPBaseClient):
 
     def __init__(self,
                  url: str,
-                 tool_call_timeout: timedelta = timedelta(seconds=5),
+                 tool_call_timeout: timedelta = timedelta(seconds=60),
+                 auth_flow_timeout: timedelta = timedelta(seconds=300),
                  reconnect_enabled: bool = True,
                  reconnect_max_attempts: int = 2,
                  reconnect_initial_backoff: float = 0.5,
                  reconnect_max_backoff: float = 50.0):
         super().__init__("sse",
                          tool_call_timeout=tool_call_timeout,
+                         auth_flow_timeout=auth_flow_timeout,
                          reconnect_enabled=reconnect_enabled,
                          reconnect_max_attempts=reconnect_max_attempts,
                          reconnect_initial_backoff=reconnect_initial_backoff,
@@ -383,13 +468,15 @@ class MCPStdioClient(MCPBaseClient):
                  command: str,
                  args: list[str] | None = None,
                  env: dict[str, str] | None = None,
-                 tool_call_timeout: timedelta = timedelta(seconds=5),
+                 tool_call_timeout: timedelta = timedelta(seconds=60),
+                 auth_flow_timeout: timedelta = timedelta(seconds=300),
                  reconnect_enabled: bool = True,
                  reconnect_max_attempts: int = 2,
                  reconnect_initial_backoff: float = 0.5,
                  reconnect_max_backoff: float = 50.0):
         super().__init__("stdio",
                          tool_call_timeout=tool_call_timeout,
+                         auth_flow_timeout=auth_flow_timeout,
                          reconnect_enabled=reconnect_enabled,
                          reconnect_max_attempts=reconnect_max_attempts,
                          reconnect_initial_backoff=reconnect_initial_backoff,
@@ -440,14 +527,18 @@ class MCPStreamableHTTPClient(MCPBaseClient):
     def __init__(self,
                  url: str,
                  auth_provider: AuthProviderBase | None = None,
-                 tool_call_timeout: timedelta = timedelta(seconds=5),
+                 user_id: str | None = None,
+                 tool_call_timeout: timedelta = timedelta(seconds=60),
+                 auth_flow_timeout: timedelta = timedelta(seconds=300),
                  reconnect_enabled: bool = True,
                  reconnect_max_attempts: int = 2,
                  reconnect_initial_backoff: float = 0.5,
                  reconnect_max_backoff: float = 50.0):
         super().__init__("streamable-http",
                          auth_provider=auth_provider,
+                         user_id=user_id,
                          tool_call_timeout=tool_call_timeout,
+                         auth_flow_timeout=auth_flow_timeout,
                          reconnect_enabled=reconnect_enabled,
                          reconnect_max_attempts=reconnect_max_attempts,
                          reconnect_initial_backoff=reconnect_initial_backoff,
@@ -490,17 +581,15 @@ class MCPToolClient:
 
     def __init__(self,
                  session: ClientSession,
-                 parent_client: "MCPBaseClient",
+                 parent_client: MCPBaseClient,
                  tool_name: str,
                  tool_description: str | None,
-                 tool_input_schema: dict | None = None,
-                 tool_call_timeout: timedelta = timedelta(seconds=5)):
+                 tool_input_schema: dict | None = None):
         self._session = session
         self._tool_name = tool_name
         self._tool_description = tool_description
         self._input_schema = (model_from_mcp_schema(self._tool_name, tool_input_schema) if tool_input_schema else None)
         self._parent_client = parent_client
-        self._tool_call_timeout = tool_call_timeout
 
         if self._parent_client is None:
             raise RuntimeError("MCPToolClient initialized without a parent client.")
@@ -535,12 +624,17 @@ class MCPToolClient:
     async def acall(self, tool_args: dict) -> str:
         """
         Call the MCP tool with the provided arguments.
+        Session context is now handled at the client level, eliminating the need for metadata injection.
 
         Args:
             tool_args (dict[str, Any]): A dictionary of key value pairs to serve as inputs for the MCP tool.
         """
-        logger.info("Calling tool %s with arguments %s", self._tool_name, tool_args)
+        if self._session is None:
+            raise RuntimeError("No session available for tool call")
+
         try:
+            # Simple tool call - session context is already in the client instance
+            logger.info("Calling tool %s", self._tool_name)
             result = await self._parent_client.call_tool(self._tool_name, tool_args)
 
             output = []
@@ -558,6 +652,6 @@ class MCPToolClient:
 
         except MCPError as e:
             format_mcp_error(e, include_traceback=False)
-            result_str = "MCPToolClient tool call failed: %s" % e.original_exception
+            result_str = f"MCPToolClient tool call failed: {e.original_exception}"
 
         return result_str