chuk-tool-processor 0.2__py3-none-any.whl → 0.4__py3-none-any.whl

chuk_tool_processor/execution/tool_executor.py

@@ -1,10 +1,12 @@
 #!/usr/bin/env python
 # chuk_tool_processor/execution/tool_executor.py
 """
-Modified ToolExecutor with true streaming support and duplicate prevention.
+Modified ToolExecutor with true streaming support and proper timeout handling.
 
 This version accesses streaming tools' stream_execute method directly
 to enable true item-by-item streaming behavior, while preventing duplicates.
+
+FIXED: Proper timeout precedence - respects strategy's default_timeout when available.
 """
 import asyncio
 from datetime import datetime, timezone
@@ -25,12 +27,14 @@ class ToolExecutor:
     
     This class provides a unified interface for executing tools using different
     execution strategies, with special support for streaming tools.
+    
+    FIXED: Proper timeout handling that respects strategy's default_timeout.
     """
 
     def __init__(
         self,
         registry: Optional[ToolRegistryInterface] = None,
-        default_timeout: float = 10.0,
+        default_timeout: Optional[float] = None,  # Made optional to allow strategy precedence
         strategy: Optional[ExecutionStrategy] = None,
         strategy_kwargs: Optional[Dict[str, Any]] = None,
     ) -> None:
@@ -39,12 +43,12 @@ class ToolExecutor:
         
         Args:
             registry: Tool registry to use for tool lookups
-            default_timeout: Default timeout for tool execution
+            default_timeout: Default timeout for tool execution (optional)
+                           If None, will use strategy's default_timeout if available
             strategy: Optional execution strategy (default: InProcessStrategy)
             strategy_kwargs: Additional arguments for the strategy constructor
         """
         self.registry = registry
-        self.default_timeout = default_timeout
         
         # Create strategy if not provided
         if strategy is None:
@@ -55,13 +59,31 @@ class ToolExecutor:
                 raise ValueError("Registry must be provided if strategy is not")
                 
             strategy_kwargs = strategy_kwargs or {}
+            
+            # If no default_timeout specified, use a reasonable default for the strategy
+            strategy_timeout = default_timeout if default_timeout is not None else 30.0
+            
             strategy = _inprocess_mod.InProcessStrategy(
                 registry,
-                default_timeout=default_timeout,
+                default_timeout=strategy_timeout,
                 **strategy_kwargs,
             )
             
         self.strategy = strategy
+        
+        # Set default timeout with proper precedence:
+        # 1. Explicit default_timeout parameter
+        # 2. Strategy's default_timeout (if available and not None)
+        # 3. Fallback to 30.0 seconds
+        if default_timeout is not None:
+            self.default_timeout = default_timeout
+            logger.debug(f"Using explicit default_timeout: {self.default_timeout}s")
+        elif hasattr(strategy, 'default_timeout') and strategy.default_timeout is not None:
+            self.default_timeout = strategy.default_timeout
+            logger.debug(f"Using strategy's default_timeout: {self.default_timeout}s")
+        else:
+            self.default_timeout = 30.0  # Conservative fallback
+            logger.debug(f"Using fallback default_timeout: {self.default_timeout}s")
 
     @property
     def supports_streaming(self) -> bool:
@@ -79,7 +101,7 @@ class ToolExecutor:
         
         Args:
             calls: List of tool calls to execute
-            timeout: Optional timeout for execution (overrides default_timeout)
+            timeout: Optional timeout for execution (overrides all defaults)
             use_cache: Whether to use cached results (for caching wrappers)
             
         Returns:
@@ -88,10 +110,13 @@ class ToolExecutor:
         if not calls:
             return []
             
-        # Use the provided timeout or fall back to default
+        # Timeout precedence:
+        # 1. Explicit timeout parameter (highest priority)
+        # 2. Executor's default_timeout (which already considers strategy's timeout)
         effective_timeout = timeout if timeout is not None else self.default_timeout
         
-        logger.debug(f"Executing {len(calls)} tool calls with timeout {effective_timeout}s")
+        logger.debug(f"Executing {len(calls)} tool calls with timeout {effective_timeout}s "
+                    f"(explicit: {timeout is not None})")
         
         # Delegate to the strategy
         return await self.strategy.run(calls, timeout=effective_timeout)
@@ -118,9 +143,12 @@ class ToolExecutor:
         if not calls:
             return
             
-        # Use the provided timeout or fall back to default
+        # Use the same timeout precedence as execute()
         effective_timeout = timeout if timeout is not None else self.default_timeout
         
+        logger.debug(f"Stream executing {len(calls)} tool calls with timeout {effective_timeout}s "
+                    f"(explicit: {timeout is not None})")
+        
         # There are two possible ways to handle streaming:
         # 1. Use the strategy's stream_run if available
         # 2. Use direct streaming for streaming tools
@@ -232,6 +260,8 @@ class ToolExecutor:
         machine = "direct-stream"
         pid = 0
         
+        logger.debug(f"Direct streaming {call.tool} with timeout {timeout}s")
+        
         # Create streaming task with timeout
         async def stream_with_timeout():
             try:
@@ -265,11 +295,16 @@ class ToolExecutor:
         try:
             if timeout:
                 await asyncio.wait_for(stream_with_timeout(), timeout)
+                logger.debug(f"Direct streaming {call.tool} completed within {timeout}s")
             else:
                 await stream_with_timeout()
+                logger.debug(f"Direct streaming {call.tool} completed (no timeout)")
         except asyncio.TimeoutError:
             # Handle timeout
             end_time = datetime.now(timezone.utc)
+            actual_duration = (end_time - start_time).total_seconds()
+            logger.debug(f"Direct streaming {call.tool} timed out after {actual_duration:.3f}s (limit: {timeout}s)")
+            
             timeout_result = ToolResult(
                 tool=call.tool,
                 result=None,
@@ -283,6 +318,8 @@ class ToolExecutor:
         except Exception as e:
             # Handle other errors
             end_time = datetime.now(timezone.utc)
+            logger.exception(f"Error in direct streaming {call.tool}: {e}")
+            
             error_result = ToolResult(
                 tool=call.tool,
                 result=None,
@@ -300,5 +337,6 @@ class ToolExecutor:
         
         This should be called during application shutdown to ensure proper cleanup.
         """
+        logger.debug("Shutting down ToolExecutor")
         if hasattr(self.strategy, "shutdown") and callable(self.strategy.shutdown):
             await self.strategy.shutdown()

chuk_tool_processor/mcp/mcp_tool.py

@@ -36,9 +36,11 @@ class MCPTool:
         servers: Optional[List[str]] = None,
         server_names: Optional[Dict[int, str]] = None,
         namespace: str = "stdio",
+        default_timeout: Optional[float] = None,  # Add default timeout support
     ) -> None:
         self.tool_name = tool_name
         self._sm: Optional[StreamManager] = stream_manager
+        self.default_timeout = default_timeout or 30.0  # Default to 30s if not specified
 
         # Boot-strap parameters (only needed if _sm is None)
         self._cfg_file = cfg_file
@@ -78,21 +80,56 @@ class MCPTool:
         return self._sm  # type: ignore[return-value]
 
     # ------------------------------------------------------------------ #
-    async def execute(self, **kwargs: Any) -> Any:
+    async def execute(self, timeout: Optional[float] = None, **kwargs: Any) -> Any:
         """
-        Forward the call to the remote MCP tool.
+        Forward the call to the remote MCP tool with timeout support.
+
+        Args:
+            timeout: Optional timeout for this specific call. If not provided,
+                    uses the instance's default_timeout.
+            **kwargs: Arguments to pass to the MCP tool.
+
+        Returns:
+            The result from the MCP tool call.
 
         Raises
         ------
         RuntimeError
             If the server returns an error payload.
+        asyncio.TimeoutError
+            If the call times out.
         """
         sm = await self._ensure_stream_manager()
-        result = await sm.call_tool(tool_name=self.tool_name, arguments=kwargs)
+        
+        # Use provided timeout, fall back to instance default, then global default
+        effective_timeout = timeout if timeout is not None else self.default_timeout
+        
+        logger.debug("Calling MCP tool '%s' with timeout: %ss", self.tool_name, effective_timeout)
+        
+        try:
+            # Pass timeout directly to StreamManager instead of wrapping with wait_for
+            result = await sm.call_tool(
+                tool_name=self.tool_name, 
+                arguments=kwargs,
+                timeout=effective_timeout
+            )
+            
+            if result.get("isError"):
+                err = result.get("error", "Unknown error")
+                logger.error("Remote MCP error from '%s': %s", self.tool_name, err)
+                raise RuntimeError(err)
+
+            return result.get("content")
+            
+        except asyncio.TimeoutError:
+            logger.warning("MCP tool '%s' timed out after %ss", self.tool_name, effective_timeout)
+            raise
+        except Exception as e:
+            logger.error("Error calling MCP tool '%s': %s", self.tool_name, e)
+            raise
 
-        if result.get("isError"):
-            err = result.get("error", "Unknown error")
-            logger.error("Remote MCP error from '%s': %s", self.tool_name, err)
-            raise RuntimeError(err)
-
-        return result.get("content")
+    # ------------------------------------------------------------------ #
+    # Legacy method name support
+    async def _aexecute(self, timeout: Optional[float] = None, **kwargs: Any) -> Any:
+        """Legacy alias for execute() method."""
+        return await self.execute(timeout=timeout, **kwargs)

chuk_tool_processor/mcp/setup_mcp_sse.py

@@ -14,6 +14,7 @@ Utility that wires up:
 
 from __future__ import annotations
 
+import os
 from typing import Dict, List, Optional, Tuple
 
 from chuk_tool_processor.core.processor import ToolProcessor
@@ -47,7 +48,26 @@ async def setup_mcp_sse(  # noqa: C901 – long, but just a config wrapper
     and return a ready-to-go :class:`ToolProcessor`.
 
     Everything is **async-native** – call with ``await``.
+    
+    NEW: Automatically detects and adds bearer token from MCP_BEARER_TOKEN
+    environment variable if not explicitly provided in server config.
     """
+    
+    # NEW: Auto-detect and add bearer token to servers if available
+    bearer_token = os.getenv("MCP_BEARER_TOKEN")
+    if bearer_token:
+        logger.info("Found MCP_BEARER_TOKEN environment variable, adding to server configs")
+        
+        # Add api_key to servers that don't already have it
+        enhanced_servers = []
+        for server in servers:
+            enhanced_server = dict(server)  # Make a copy
+            if "api_key" not in enhanced_server and bearer_token:
+                enhanced_server["api_key"] = bearer_token
+                logger.info("Added bearer token to server: %s", enhanced_server.get("name", "unnamed"))
+            enhanced_servers.append(enhanced_server)
+        servers = enhanced_servers
+    
     # 1️⃣  connect to the remote MCP servers
     stream_manager = await StreamManager.create_with_sse(
         servers=servers,
@@ -76,4 +96,4 @@ async def setup_mcp_sse(  # noqa: C901 – long, but just a config wrapper
         "" if len(registered) == 1 else "s",
         namespace,
     )
-    return processor, stream_manager
+    return processor, stream_manager

chuk_tool_processor/mcp/stream_manager.py

@@ -47,9 +47,16 @@ class StreamManager:
         servers: List[str],
         server_names: Optional[Dict[int, str]] = None,
         transport_type: str = "stdio",
+        default_timeout: float = 30.0,  # ADD: For consistency
     ) -> "StreamManager":
         inst = cls()
-        await inst.initialize(config_file, servers, server_names, transport_type)
+        await inst.initialize(
+            config_file, 
+            servers, 
+            server_names, 
+            transport_type,
+            default_timeout=default_timeout  # PASS THROUGH
+        )
         return inst
 
     @classmethod
@@ -57,9 +64,16 @@ class StreamManager:
         cls,
         servers: List[Dict[str, str]],
         server_names: Optional[Dict[int, str]] = None,
+        connection_timeout: float = 10.0,  # ADD: For SSE connection setup
+        default_timeout: float = 30.0,     # ADD: For tool execution
     ) -> "StreamManager":
         inst = cls()
-        await inst.initialize_with_sse(servers, server_names)
+        await inst.initialize_with_sse(
+            servers, 
+            server_names,
+            connection_timeout=connection_timeout,  # PASS THROUGH
+            default_timeout=default_timeout         # PASS THROUGH
+        )
         return inst
 
     # ------------------------------------------------------------------ #
@@ -71,6 +85,7 @@ class StreamManager:
         servers: List[str],
         server_names: Optional[Dict[int, str]] = None,
         transport_type: str = "stdio",
+        default_timeout: float = 30.0,  # ADD: For consistency
     ) -> None:
         async with self._lock:
             self.server_names = server_names or {}
@@ -81,7 +96,24 @@ class StreamManager:
                         params = await load_config(config_file, server_name)
                         transport: MCPBaseTransport = StdioTransport(params)
                     elif transport_type == "sse":
-                        transport = SSETransport("http://localhost:8000")
+                        # WARNING: For SSE transport, prefer using create_with_sse() instead
+                        # This is a fallback for backward compatibility
+                        logger.warning("Using SSE transport in initialize() - consider using initialize_with_sse() instead")
+                        
+                        # Try to extract URL from params or use localhost as fallback
+                        if isinstance(params, dict) and 'url' in params:
+                            sse_url = params['url']
+                            api_key = params.get('api_key')
+                        else:
+                            sse_url = "http://localhost:8000"
+                            api_key = None
+                            logger.warning(f"No URL configured for SSE transport, using default: {sse_url}")
+                        
+                        transport = SSETransport(
+                            sse_url,
+                            api_key,
+                            default_timeout=default_timeout
+                        )
                     else:
                         logger.error("Unsupported transport type: %s", transport_type)
                         continue
@@ -125,6 +157,8 @@ class StreamManager:
         self,
         servers: List[Dict[str, str]],
         server_names: Optional[Dict[int, str]] = None,
+        connection_timeout: float = 10.0,  # ADD: For SSE connection setup
+        default_timeout: float = 30.0,     # ADD: For tool execution
     ) -> None:
         async with self._lock:
             self.server_names = server_names or {}
@@ -135,7 +169,14 @@ class StreamManager:
                     logger.error("Bad server config: %s", cfg)
                     continue
                 try:
-                    transport = SSETransport(url, cfg.get("api_key"))
+                    # FIXED: Pass timeout parameters to SSETransport
+                    transport = SSETransport(
+                        url, 
+                        cfg.get("api_key"),
+                        connection_timeout=connection_timeout,  # ADD THIS
+                        default_timeout=default_timeout         # ADD THIS
+                    )
+                    
                     if not await transport.initialize():
                         logger.error("Failed to init SSE %s", name)
                         continue
@@ -265,7 +306,20 @@ class StreamManager:
         tool_name: str,
         arguments: Dict[str, Any],
         server_name: Optional[str] = None,
+        timeout: Optional[float] = None,  # Timeout parameter already exists
     ) -> Dict[str, Any]:
+        """
+        Call a tool on the appropriate server with timeout support.
+        
+        Args:
+            tool_name: Name of the tool to call
+            arguments: Arguments to pass to the tool
+            server_name: Optional server name (auto-detected if not provided)
+            timeout: Optional timeout for the call
+            
+        Returns:
+            Dictionary containing the tool result or error
+        """
         server_name = server_name or self.get_server_for_tool(tool_name)
         if not server_name or server_name not in self.transports:
             # wording kept exactly for unit-test expectation
@@ -273,8 +327,42 @@ class StreamManager:
                 "isError": True,
                 "error": f"No server found for tool: {tool_name}",
             }
-        return await self.transports[server_name].call_tool(tool_name, arguments)
-
+        
+        transport = self.transports[server_name]
+        
+        # Apply timeout if specified
+        if timeout is not None:
+            logger.debug("Calling tool '%s' with %ss timeout", tool_name, timeout)
+            try:
+                # ENHANCED: Pass timeout to transport.call_tool if it supports it
+                if hasattr(transport, 'call_tool'):
+                    import inspect
+                    sig = inspect.signature(transport.call_tool)
+                    if 'timeout' in sig.parameters:
+                        # Transport supports timeout parameter - pass it through
+                        return await transport.call_tool(tool_name, arguments, timeout=timeout)
+                    else:
+                        # Transport doesn't support timeout - use asyncio.wait_for wrapper
+                        return await asyncio.wait_for(
+                            transport.call_tool(tool_name, arguments),
+                            timeout=timeout
+                        )
+                else:
+                    # Fallback to asyncio.wait_for
+                    return await asyncio.wait_for(
+                        transport.call_tool(tool_name, arguments),
+                        timeout=timeout
+                    )
+            except asyncio.TimeoutError:
+                logger.warning("Tool '%s' timed out after %ss", tool_name, timeout)
+                return {
+                    "isError": True,
+                    "error": f"Tool call timed out after {timeout}s",
+                }
+        else:
+            # No timeout specified, call directly
+            return await transport.call_tool(tool_name, arguments)
+        
     # ------------------------------------------------------------------ #
     #  shutdown                                                          #
     # ------------------------------------------------------------------ #
@@ -318,4 +406,4 @@ class StreamManager:
     # convenience alias
     @property
     def streams(self) -> List[Tuple[Any, Any]]:  # pragma: no cover
-        return self.get_streams()
+        return self.get_streams()

chuk_tool_processor/mcp/transport/sse_transport.py

@@ -8,12 +8,15 @@ This transport:
 3. Sends MCP initialize handshake FIRST
 4. Only then proceeds with tools/list and tool calls
 5. Handles async responses via SSE message events
+
+FIXED: All hardcoded timeouts are now configurable parameters.
 """
 from __future__ import annotations
 
 import asyncio
 import contextlib
 import json
+import os
 from typing import Any, Dict, List, Optional
 
 import httpx
@@ -23,7 +26,8 @@ from .base_transport import MCPBaseTransport
 # --------------------------------------------------------------------------- #
 # Helpers                                                                     #
 # --------------------------------------------------------------------------- #
-DEFAULT_TIMEOUT = 30.0  # Longer timeout for real servers
+DEFAULT_TIMEOUT = 30.0  # Default timeout for tool calls
+DEFAULT_CONNECTION_TIMEOUT = 10.0  # Default timeout for connection setup
 HEADERS_JSON: Dict[str, str] = {"accept": "application/json"}
 
 
@@ -46,9 +50,33 @@ class SSETransport(MCPBaseTransport):
     5. Waits for async responses via SSE message events
     """
 
-    def __init__(self, url: str, api_key: Optional[str] = None) -> None:
+    def __init__(
+        self, 
+        url: str, 
+        api_key: Optional[str] = None,
+        connection_timeout: float = DEFAULT_CONNECTION_TIMEOUT,
+        default_timeout: float = DEFAULT_TIMEOUT
+    ) -> None:
+        """
+        Initialize SSE Transport with configurable timeouts.
+        
+        Args:
+            url: Base URL for the MCP server
+            api_key: Optional API key for authentication
+            connection_timeout: Timeout for connection setup (default: 10.0s)
+            default_timeout: Default timeout for tool calls (default: 30.0s)
+        """
         self.base_url = url.rstrip("/")
         self.api_key = api_key
+        self.connection_timeout = connection_timeout
+        self.default_timeout = default_timeout
+        
+        # NEW: Auto-detect bearer token from environment if not provided
+        if not self.api_key:
+            bearer_token = os.getenv("MCP_BEARER_TOKEN")
+            if bearer_token:
+                self.api_key = bearer_token
+                print(f"🔑 Using bearer token from MCP_BEARER_TOKEN environment variable")
 
         # httpx client (None until initialise)
         self._client: httpx.AsyncClient | None = None
@@ -75,11 +103,16 @@ class SSETransport(MCPBaseTransport):
 
         headers = {}
         if self.api_key:
-            headers["authorization"] = self.api_key
+            # NEW: Handle both "Bearer token" and just "token" formats
+            if self.api_key.startswith("Bearer "):
+                headers["Authorization"] = self.api_key
+            else:
+                headers["Authorization"] = f"Bearer {self.api_key}"
+            print(f"🔑 Added Authorization header to httpx client")
 
         self._client = httpx.AsyncClient(
             headers=headers,
-            timeout=DEFAULT_TIMEOUT,
+            timeout=self.default_timeout,  # Use configurable timeout
         )
         self.session = self._client
 
@@ -87,8 +120,8 @@ class SSETransport(MCPBaseTransport):
         self._sse_task = asyncio.create_task(self._handle_sse_connection())
         
         try:
-            # Wait for endpoint event (up to 10 seconds)
-            await asyncio.wait_for(self._connected.wait(), timeout=10.0)
+            # FIXED: Use configurable connection timeout instead of hardcoded 10.0
+            await asyncio.wait_for(self._connected.wait(), timeout=self.connection_timeout)
             
             # NEW: Send MCP initialize handshake
             if await self._initialize_mcp_session():
@@ -223,6 +256,11 @@ class SSETransport(MCPBaseTransport):
                         
                         if event_type == "endpoint":
                             # Got the endpoint URL for messages - construct full URL
+                            # NEW: Handle URLs that need trailing slash fix
+                            if "/messages?" in data and "/messages/?" not in data:
+                                data = data.replace("/messages?", "/messages/?", 1)
+                                print(f"🔧 Fixed URL redirect: added trailing slash")
+                            
                             self._message_url = f"{self.base_url}{data}"
                             
                             # Extract session_id if present
@@ -267,7 +305,8 @@ class SSETransport(MCPBaseTransport):
         if not self._initialized.is_set():
             print("⏳ Waiting for MCP initialization...")
             try:
-                await asyncio.wait_for(self._initialized.wait(), timeout=10.0)
+                # FIXED: Use configurable connection timeout instead of hardcoded 10.0
+                await asyncio.wait_for(self._initialized.wait(), timeout=self.connection_timeout)
             except asyncio.TimeoutError:
                 print("❌ Timeout waiting for MCP initialization")
                 return []
@@ -293,8 +332,23 @@ class SSETransport(MCPBaseTransport):
             
         return []
 
-    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
-        """Execute a tool call using the MCP protocol."""
+    async def call_tool(
+        self, 
+        tool_name: str, 
+        arguments: Dict[str, Any],
+        timeout: Optional[float] = None
+    ) -> Dict[str, Any]:
+        """
+        Execute a tool call using the MCP protocol.
+        
+        Args:
+            tool_name: Name of the tool to call
+            arguments: Arguments to pass to the tool
+            timeout: Optional timeout for this specific call
+            
+        Returns:
+            Dictionary containing the tool result or error
+        """
         # NEW: Ensure initialization before tool calls
         if not self._initialized.is_set():
             return {"isError": True, "error": "MCP session not initialized"}
@@ -313,7 +367,9 @@ class SSETransport(MCPBaseTransport):
                 }
             }
             
-            response = await self._send_message(message)
+            # Use provided timeout or fall back to default
+            effective_timeout = timeout if timeout is not None else self.default_timeout
+            response = await self._send_message(message, timeout=effective_timeout)
             
             # Process MCP response
             if "error" in response:
@@ -345,8 +401,21 @@ class SSETransport(MCPBaseTransport):
         except Exception as e:
             return {"isError": True, "error": str(e)}
 
-    async def _send_message(self, message: Dict[str, Any]) -> Dict[str, Any]:
-        """Send a JSON-RPC message to the server and wait for async response."""
+    async def _send_message(
+        self, 
+        message: Dict[str, Any],
+        timeout: Optional[float] = None
+    ) -> Dict[str, Any]:
+        """
+        Send a JSON-RPC message to the server and wait for async response.
+        
+        Args:
+            message: JSON-RPC message to send
+            timeout: Optional timeout for this specific message
+            
+        Returns:
+            Response message from the server
+        """
         if not self._client or not self._message_url:
             raise RuntimeError("Transport not properly initialized")
 
@@ -354,6 +423,9 @@ class SSETransport(MCPBaseTransport):
         if not message_id:
             raise ValueError("Message must have an ID")
 
+        # Use provided timeout or fall back to default
+        effective_timeout = timeout if timeout is not None else self.default_timeout
+
         # Create a future for this request
         future = asyncio.Future()
         async with self._message_lock:
@@ -373,7 +445,8 @@ class SSETransport(MCPBaseTransport):
             if response.status_code == 202:
                 # Server accepted - wait for async response via SSE
                 try:
-                    response_message = await asyncio.wait_for(future, timeout=30.0)
+                    # FIXED: Use effective_timeout instead of hardcoded 30.0
+                    response_message = await asyncio.wait_for(future, timeout=effective_timeout)
                     return response_message
                 except asyncio.TimeoutError:
                     raise RuntimeError(f"Timeout waiting for response to message {message_id}")