chuk-tool-processor 0.3__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/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,7 @@ class StreamManager:
         tool_name: str,
         arguments: Dict[str, Any],
         server_name: Optional[str] = None,
-        timeout: Optional[float] = None,  # Add timeout parameter
+        timeout: Optional[float] = None,  # Timeout parameter already exists
     ) -> Dict[str, Any]:
         """
         Call a tool on the appropriate server with timeout support.
@@ -293,10 +334,25 @@ class StreamManager:
         if timeout is not None:
             logger.debug("Calling tool '%s' with %ss timeout", tool_name, timeout)
             try:
-                return await asyncio.wait_for(
-                    transport.call_tool(tool_name, arguments),
-                    timeout=timeout
-                )
+                # 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 {
@@ -350,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,6 +8,8 @@ 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
 
@@ -24,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"}
 
 
@@ -47,9 +50,26 @@ 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:
@@ -92,7 +112,7 @@ class SSETransport(MCPBaseTransport):
 
         self._client = httpx.AsyncClient(
             headers=headers,
-            timeout=DEFAULT_TIMEOUT,
+            timeout=self.default_timeout,  # Use configurable timeout
         )
         self.session = self._client
 
@@ -100,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():
@@ -285,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 []
@@ -311,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"}
@@ -331,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:
@@ -363,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")
 
@@ -372,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:
@@ -391,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}")