chuk-tool-processor 0.6.6__py3-none-any.whl → 0.6.9__py3-none-any.whl

chuk_tool_processor/mcp/stream_manager.py

@@ -1,6 +1,6 @@
 # chuk_tool_processor/mcp/stream_manager.py
 """
-StreamManager for CHUK Tool Processor - Enhanced with robust shutdown handling
+StreamManager for CHUK Tool Processor - Enhanced with robust shutdown handling and headers support
 """
 from __future__ import annotations
 
@@ -27,12 +27,12 @@ class StreamManager:
     """
     Manager for MCP server streams with support for multiple transport types.
     
-    Enhanced with robust shutdown handling to prevent event loop closure issues.
+    Enhanced with robust shutdown handling and proper headers support.
     
     Updated to support the latest transports:
     - STDIO (process-based)
-    - SSE (Server-Sent Events) 
-    - HTTP Streamable (modern replacement for SSE, spec 2025-03-26)
+    - SSE (Server-Sent Events) with headers support
+    - HTTP Streamable (modern replacement for SSE, spec 2025-03-26) with graceful headers handling
     """
 
     def __init__(self) -> None:
@@ -175,6 +175,7 @@ class StreamManager:
         transport_type: str = "stdio",
         default_timeout: float = 30.0,
     ) -> None:
+        """Initialize with graceful headers handling for all transport types."""
         if self._closed:
             raise RuntimeError("Cannot initialize a closed StreamManager")
             
@@ -193,16 +194,24 @@ class StreamManager:
                         if isinstance(params, dict) and 'url' in params:
                             sse_url = params['url']
                             api_key = params.get('api_key')
+                            headers = params.get('headers', {})
                         else:
                             sse_url = "http://localhost:8000"
                             api_key = None
+                            headers = {}
                             logger.warning("No URL configured for SSE transport, using default: %s", sse_url)
                         
-                        transport = SSETransport(
-                            sse_url,
-                            api_key,
-                            default_timeout=default_timeout
-                        )
+                        # Build SSE transport with optional headers
+                        transport_params = {
+                            'url': sse_url,
+                            'api_key': api_key,
+                            'default_timeout': default_timeout
+                        }
+                        if headers:
+                            transport_params['headers'] = headers
+                        
+                        transport = SSETransport(**transport_params)
+                        
                     elif transport_type == "http_streamable":
                         logger.warning("Using HTTP Streamable transport in initialize() - consider using initialize_with_http_streamable() instead")
                         params = await load_config(config_file, server_name)
@@ -210,19 +219,28 @@ class StreamManager:
                         if isinstance(params, dict) and 'url' in params:
                             http_url = params['url']
                             api_key = params.get('api_key')
+                            headers = params.get('headers', {})
                             session_id = params.get('session_id')
                         else:
                             http_url = "http://localhost:8000"
                             api_key = None
+                            headers = {}
                             session_id = None
                             logger.warning("No URL configured for HTTP Streamable transport, using default: %s", http_url)
                         
-                        transport = HTTPStreamableTransport(
-                            http_url,
-                            api_key,
-                            default_timeout=default_timeout,
-                            session_id=session_id
-                        )
+                        # Build HTTP transport (headers not supported yet)
+                        transport_params = {
+                            'url': http_url,
+                            'api_key': api_key,
+                            'default_timeout': default_timeout,
+                            'session_id': session_id
+                        }
+                        # Note: headers not added until HTTPStreamableTransport supports them
+                        if headers:
+                            logger.debug("Headers provided but not supported in HTTPStreamableTransport yet")
+                        
+                        transport = HTTPStreamableTransport(**transport_params)
+                        
                     else:
                         logger.error("Unsupported transport type: %s", transport_type)
                         continue
@@ -271,6 +289,7 @@ class StreamManager:
         connection_timeout: float = 10.0,
         default_timeout: float = 30.0,
     ) -> None:
+        """Initialize with SSE transport with optional headers support."""
         if self._closed:
             raise RuntimeError("Cannot initialize a closed StreamManager")
             
@@ -283,12 +302,21 @@ class StreamManager:
                     logger.error("Bad server config: %s", cfg)
                     continue
                 try:
-                    transport = SSETransport(
-                        url, 
-                        cfg.get("api_key"),
-                        connection_timeout=connection_timeout,
-                        default_timeout=default_timeout
-                    )
+                    # Build SSE transport parameters with optional headers
+                    transport_params = {
+                        'url': url,
+                        'api_key': cfg.get("api_key"),
+                        'connection_timeout': connection_timeout,
+                        'default_timeout': default_timeout
+                    }
+                    
+                    # Add headers if provided
+                    headers = cfg.get("headers", {})
+                    if headers:
+                        logger.debug("SSE %s: Using configured headers: %s", name, list(headers.keys()))
+                        transport_params['headers'] = headers
+                    
+                    transport = SSETransport(**transport_params)
                     
                     if not await asyncio.wait_for(transport.initialize(), timeout=connection_timeout):
                         logger.error("Failed to init SSE %s", name)
@@ -326,7 +354,7 @@ class StreamManager:
         connection_timeout: float = 30.0,
         default_timeout: float = 30.0,
     ) -> None:
-        """Initialize with HTTP Streamable transport (modern MCP spec 2025-03-26)."""
+        """Initialize with HTTP Streamable transport with graceful headers handling."""
         if self._closed:
             raise RuntimeError("Cannot initialize a closed StreamManager")
             
@@ -339,13 +367,23 @@ class StreamManager:
                     logger.error("Bad server config: %s", cfg)
                     continue
                 try:
-                    transport = HTTPStreamableTransport(
-                        url, 
-                        cfg.get("api_key"),
-                        connection_timeout=connection_timeout,
-                        default_timeout=default_timeout,
-                        session_id=cfg.get("session_id")
-                    )
+                    # Build HTTP Streamable transport parameters
+                    transport_params = {
+                        'url': url,
+                        'api_key': cfg.get("api_key"),
+                        'connection_timeout': connection_timeout,
+                        'default_timeout': default_timeout,
+                        'session_id': cfg.get("session_id")
+                    }
+                    
+                    # Handle headers if provided (for future HTTPStreamableTransport support)
+                    headers = cfg.get("headers", {})
+                    if headers:
+                        logger.debug("HTTP Streamable %s: Headers provided but not yet supported in transport", name)
+                        # TODO: Add headers support when HTTPStreamableTransport is updated
+                        # transport_params['headers'] = headers
+                    
+                    transport = HTTPStreamableTransport(**transport_params)
                     
                     if not await asyncio.wait_for(transport.initialize(), timeout=connection_timeout):
                         logger.error("Failed to init HTTP Streamable %s", name)

chuk_tool_processor/mcp/transport/__init__.py

@@ -1,39 +1,22 @@
 # chuk_tool_processor/mcp/transport/__init__.py
 """
-MCP Transport module providing multiple transport implementations.
+MCP Transport module providing consistent transport implementations.
+
+All transports now follow the same interface and provide consistent behavior:
+- Standardized initialization and cleanup
+- Unified metrics and monitoring
+- Consistent error handling and timeouts
+- Shared response normalization
 """
 
 from .base_transport import MCPBaseTransport
-
-# Always available transports
-try:
-    from .stdio_transport import StdioTransport
-    HAS_STDIO_TRANSPORT = True
-except ImportError:
-    StdioTransport = None
-    HAS_STDIO_TRANSPORT = False
-
-# Conditionally available transports
-try:
-    from .sse_transport import SSETransport
-    HAS_SSE_TRANSPORT = True
-except ImportError:
-    SSETransport = None
-    HAS_SSE_TRANSPORT = False
-
-try:
-    from .http_streamable_transport import HTTPStreamableTransport
-    HAS_HTTP_STREAMABLE_TRANSPORT = True
-except ImportError:
-    HTTPStreamableTransport = None
-    HAS_HTTP_STREAMABLE_TRANSPORT = False
+from .stdio_transport import StdioTransport
+from .sse_transport import SSETransport
+from .http_streamable_transport import HTTPStreamableTransport
 
 __all__ = [
     "MCPBaseTransport",
     "StdioTransport",
     "SSETransport", 
     "HTTPStreamableTransport",
-    "HAS_STDIO_TRANSPORT",
-    "HAS_SSE_TRANSPORT",
-    "HAS_HTTP_STREAMABLE_TRANSPORT"
 ]

chuk_tool_processor/mcp/transport/base_transport.py

@@ -1,103 +1,254 @@
 # chuk_tool_processor/mcp/transport/base_transport.py
 """
-Abstract transport layer for MCP communication.
+Abstract base class for MCP transports with complete interface definition.
 """
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 
 class MCPBaseTransport(ABC):
     """
-    Abstract base class for MCP transport mechanisms.
+    Abstract base class for all MCP transport implementations.
+    
+    Defines the complete interface that all transports must implement
+    for consistency across stdio, SSE, and HTTP streamable transports.
     """
 
     # ------------------------------------------------------------------ #
-    #  connection lifecycle                                              #
+    #  Core connection lifecycle                                         #
     # ------------------------------------------------------------------ #
     @abstractmethod
     async def initialize(self) -> bool:
         """
-        Establish the connection.
-
-        Returns
-        -------
-        bool
-            ``True`` if the connection was initialised successfully.
+        Initialize the transport connection.
+        
+        Returns:
+            True if initialization was successful, False otherwise.
         """
-        raise NotImplementedError
+        pass
 
     @abstractmethod
     async def close(self) -> None:
-        """Tear down the connection and release all resources."""
-        raise NotImplementedError
+        """Close the transport connection and clean up all resources."""
+        pass
 
     # ------------------------------------------------------------------ #
-    #  diagnostics                                                       #
+    #  Health and diagnostics                                            #
     # ------------------------------------------------------------------ #
     @abstractmethod
     async def send_ping(self) -> bool:
         """
-        Send a **ping** request.
+        Send a ping to verify the connection is alive.
+        
+        Returns:
+            True if ping was successful, False otherwise.
+        """
+        pass
 
-        Returns
-        -------
-        bool
-            ``True`` on success, ``False`` otherwise.
+    @abstractmethod
+    def is_connected(self) -> bool:
+        """
+        Check if the transport is connected and ready for operations.
+        
+        Returns:
+            True if connected, False otherwise.
         """
-        raise NotImplementedError
+        pass
 
     # ------------------------------------------------------------------ #
-    #  tool handling                                                     #
+    #  Core MCP operations                                               #
     # ------------------------------------------------------------------ #
     @abstractmethod
     async def get_tools(self) -> List[Dict[str, Any]]:
         """
-        Return a list with *all* tool definitions exposed by the server.
+        Get the list of available tools from the server.
+        
+        Returns:
+            List of tool definitions.
         """
-        raise NotImplementedError
+        pass
 
     @abstractmethod
-    async def call_tool(
-        self, tool_name: str, arguments: Dict[str, Any]
-    ) -> Dict[str, Any]:
+    async def call_tool(self, tool_name: str, arguments: Dict[str, Any], 
+                       timeout: Optional[float] = None) -> Dict[str, Any]:
         """
-        Execute *tool_name* with *arguments* and return the normalised result.
+        Call a tool with the given arguments.
+        
+        Args:
+            tool_name: Name of the tool to call
+            arguments: Arguments to pass to the tool
+            timeout: Optional timeout for the operation
+            
+        Returns:
+            Dictionary with 'isError' boolean and either 'content' or 'error'
         """
-        raise NotImplementedError
+        pass
 
-    # ------------------------------------------------------------------ #
-    #  new: resources & prompts                                          #
-    # ------------------------------------------------------------------ #
     @abstractmethod
     async def list_resources(self) -> Dict[str, Any]:
         """
-        Retrieve the server's resources catalogue.
-
-        Expected shape::
-            { "resources": [ {...}, ... ], "nextCursor": "…", … }
+        List available resources from the server.
+        
+        Returns:
+            Dictionary containing resources list or empty dict if not supported.
         """
-        raise NotImplementedError
+        pass
 
     @abstractmethod
     async def list_prompts(self) -> Dict[str, Any]:
         """
-        Retrieve the server's prompt catalogue.
-
-        Expected shape::
-            { "prompts": [ {...}, ... ], "nextCursor": "…", … }
+        List available prompts from the server.
+        
+        Returns:
+            Dictionary containing prompts list or empty dict if not supported.
         """
-        raise NotImplementedError
+        pass
 
     # ------------------------------------------------------------------ #
-    #  optional helper (non-abstract)                                    #
+    #  Metrics and monitoring (all transports should support these)     #
     # ------------------------------------------------------------------ #
-    def get_streams(self):
+    @abstractmethod
+    def get_metrics(self) -> Dict[str, Any]:
+        """
+        Get performance and connection metrics.
+        
+        Returns:
+            Dictionary containing metrics data.
         """
-        Return a list of ``(read_stream, write_stream)`` tuples.
+        pass
+
+    @abstractmethod
+    def reset_metrics(self) -> None:
+        """Reset performance metrics to initial state."""
+        pass
 
-        Transports that do not expose their low-level streams can simply leave
-        the default implementation (which returns an empty list).
+    # ------------------------------------------------------------------ #
+    #  Backward compatibility and utility methods                       #
+    # ------------------------------------------------------------------ #
+    def get_streams(self) -> List[tuple]:
+        """
+        Get underlying stream objects for backward compatibility.
+        
+        Returns:
+            List of (read_stream, write_stream) tuples, empty if not applicable.
         """
         return []
+
+    # ------------------------------------------------------------------ #
+    #  Context manager support (all transports should support this)     #
+    # ------------------------------------------------------------------ #
+    async def __aenter__(self):
+        """Context manager entry."""
+        success = await self.initialize()
+        if not success:
+            raise RuntimeError(f"Failed to initialize {self.__class__.__name__}")
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Context manager cleanup."""
+        await self.close()
+
+    # ------------------------------------------------------------------ #
+    #  Shared helper methods for response normalization                  #
+    # ------------------------------------------------------------------ #
+    def _normalize_mcp_response(self, response: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Normalize MCP response to consistent format.
+        
+        This provides shared logic for all transports to ensure consistent
+        response format regardless of transport type.
+        """
+        # Handle explicit error in response
+        if "error" in response:
+            error_info = response["error"]
+            if isinstance(error_info, dict):
+                error_msg = error_info.get("message", "Unknown error")
+            else:
+                error_msg = str(error_info)
+            
+            return {
+                "isError": True,
+                "error": error_msg
+            }
+
+        # Handle successful response with result
+        if "result" in response:
+            result = response["result"]
+            
+            if isinstance(result, dict) and "content" in result:
+                return {
+                    "isError": False,
+                    "content": self._extract_mcp_content(result["content"])
+                }
+            else:
+                return {
+                    "isError": False,
+                    "content": result
+                }
+
+        # Handle direct content-based response
+        if "content" in response:
+            return {
+                "isError": False,
+                "content": self._extract_mcp_content(response["content"])
+            }
+
+        # Fallback
+        return {
+            "isError": False,
+            "content": response
+        }
+
+    def _extract_mcp_content(self, content_list: Any) -> Any:
+        """
+        Extract content from MCP content format.
+        
+        Handles the standard MCP content format where content is a list
+        of content items with type and data.
+        """
+        if not isinstance(content_list, list) or not content_list:
+            return content_list
+        
+        # Handle single content item
+        if len(content_list) == 1:
+            content_item = content_list[0]
+            if isinstance(content_item, dict):
+                if content_item.get("type") == "text":
+                    text_content = content_item.get("text", "")
+                    # Try to parse JSON, fall back to plain text
+                    try:
+                        import json
+                        return json.loads(text_content)
+                    except json.JSONDecodeError:
+                        return text_content
+                else:
+                    return content_item
+        
+        # Multiple content items - return as-is
+        return content_list
+
+    # ------------------------------------------------------------------ #
+    #  Standard string representation                                    #
+    # ------------------------------------------------------------------ #
+    def __repr__(self) -> str:
+        """Standard string representation for all transports."""
+        status = "initialized" if getattr(self, '_initialized', False) else "not initialized"
+        
+        # Add metrics info if available
+        metrics_info = ""
+        if hasattr(self, 'enable_metrics') and getattr(self, 'enable_metrics', False):
+            metrics = self.get_metrics()
+            if metrics.get("total_calls", 0) > 0:
+                success_rate = (metrics.get("successful_calls", 0) / metrics["total_calls"]) * 100
+                metrics_info = f", calls: {metrics['total_calls']}, success: {success_rate:.1f}%"
+        
+        # Add transport-specific info - FIXED FORMAT
+        transport_info = ""
+        if hasattr(self, 'url'):
+            transport_info = f", url={self.url}"  # Fixed: was "url: "
+        elif hasattr(self, 'server_params') and hasattr(self.server_params, 'command'):
+            transport_info = f", command={self.server_params.command}"  # Fixed: was "command: "
+        
+        return f"{self.__class__.__name__}(status={status}{transport_info}{metrics_info})"