chuk-tool-processor 0.6.7__py3-none-any.whl → 0.6.10__py3-none-any.whl

chuk_tool_processor/mcp/transport/sse_transport.py

@@ -1,14 +1,24 @@
 # chuk_tool_processor/mcp/transport/sse_transport.py
 """
-Fixed SSE transport that matches your server's actual behavior.
-Based on your working debug script with smart URL detection.
+SSE transport for MCP communication.
+
+Implements Server-Sent Events transport with two-step async pattern:
+1. POST messages to /messages endpoint
+2. Receive responses via SSE stream
+
+Note: This transport is deprecated in favor of HTTP Streamable (spec 2025-03-26)
+but remains supported for backward compatibility.
+
+FIXED: Updated to support both old format (/messages/) and new event-based format
+(event: endpoint + data: https://...) for session discovery.
 """
 from __future__ import annotations
 
 import asyncio
 import json
+import time
 import uuid
-from typing import Dict, Any, List, Optional, Tuple
+from typing import Dict, Any, List, Optional
 import logging
 
 import httpx
@@ -20,28 +30,45 @@ logger = logging.getLogger(__name__)
 
 class SSETransport(MCPBaseTransport):
     """
-    SSE transport that works with your server's two-step async pattern:
-    1. POST messages to /messages endpoint
-    2. Receive responses via SSE stream
+    SSE transport implementing the MCP protocol over Server-Sent Events.
+    
+    This transport uses a dual-connection approach:
+    - SSE stream for receiving responses
+    - HTTP POST for sending requests
+    
+    FIXED: Supports both old and new session discovery formats.
     """
 
     def __init__(self, url: str, api_key: Optional[str] = None, 
                  headers: Optional[Dict[str, str]] = None,
-                 connection_timeout: float = 30.0, default_timeout: float = 30.0):
-        """Initialize SSE transport."""
+                 connection_timeout: float = 30.0, 
+                 default_timeout: float = 30.0,
+                 enable_metrics: bool = True):
+        """
+        Initialize SSE transport.
+        
+        Args:
+            url: Base URL for the MCP server
+            api_key: Optional API key for authentication
+            headers: Optional custom headers
+            connection_timeout: Timeout for initial connection setup
+            default_timeout: Default timeout for operations
+            enable_metrics: Whether to track performance metrics
+        """
         self.url = url.rstrip('/')
         self.api_key = api_key
         self.configured_headers = headers or {}
         self.connection_timeout = connection_timeout
         self.default_timeout = default_timeout
+        self.enable_metrics = enable_metrics
         
-        # DEBUG: Log what we received
-        logger.debug("SSE Transport initialized with:")
-        logger.debug("  URL: %s", self.url)
-        logger.debug("  API Key: %s", "***" if api_key else None)
-        logger.debug("  Headers: %s", {k: v[:10] + "..." if len(v) > 10 else v for k, v in self.configured_headers.items()})
+        logger.debug("SSE Transport initialized with URL: %s", self.url)
+        if self.api_key:
+            logger.debug("API key configured for authentication")
+        if self.configured_headers:
+            logger.debug("Custom headers configured: %s", list(self.configured_headers.keys()))
         
-        # State
+        # Connection state
         self.session_id = None
         self.message_url = None
         self.pending_requests: Dict[str, asyncio.Future] = {}
@@ -51,10 +78,23 @@ class SSETransport(MCPBaseTransport):
         self.stream_client = None
         self.send_client = None
         
-        # SSE stream
+        # SSE stream management
         self.sse_task = None
         self.sse_response = None
         self.sse_stream_context = None
+        
+        # Performance metrics (consistent with other transports)
+        self._metrics = {
+            "total_calls": 0,
+            "successful_calls": 0,
+            "failed_calls": 0,
+            "total_time": 0.0,
+            "avg_response_time": 0.0,
+            "last_ping_time": None,
+            "initialization_time": None,
+            "session_discoveries": 0,
+            "stream_errors": 0
+        }
 
     def _construct_sse_url(self, base_url: str) -> str:
         """
@@ -62,53 +102,48 @@ class SSETransport(MCPBaseTransport):
         
         Smart detection to avoid double-appending /sse if already present.
         """
-        # Remove trailing slashes
         base_url = base_url.rstrip('/')
         
-        # Check if URL already ends with /sse
         if base_url.endswith('/sse'):
-            # Already has /sse, use as-is
             logger.debug("URL already contains /sse endpoint: %s", base_url)
             return base_url
         
-        # Append /sse to the base URL
         sse_url = f"{base_url}/sse"
-        logger.debug("Appending /sse to base URL: %s -> %s", base_url, sse_url)
+        logger.debug("Constructed SSE URL: %s -> %s", base_url, sse_url)
         return sse_url
 
     def _get_headers(self) -> Dict[str, str]:
-        """Get headers with auth if available."""
+        """Get headers with authentication and custom headers."""
         headers = {}
         
         # Add configured headers first
         if self.configured_headers:
             headers.update(self.configured_headers)
         
-        # Add API key as Bearer token if provided (this will override any Authorization header)
+        # Add API key as Bearer token if provided (overrides Authorization header)
         if self.api_key:
             headers['Authorization'] = f'Bearer {self.api_key}'
         
-        # DEBUG: Log what headers we're sending
-        logger.debug("Sending headers: %s", {k: v[:10] + "..." if len(v) > 10 else v for k, v in headers.items()})
-        
         return headers
 
     async def initialize(self) -> bool:
-        """Initialize SSE connection and MCP handshake."""
+        """Initialize SSE connection and perform MCP handshake."""
         if self._initialized:
             logger.warning("Transport already initialized")
             return True
         
+        start_time = time.time()
+        
         try:
             logger.debug("Initializing SSE transport...")
             
-            # Create HTTP clients
+            # Create HTTP clients with appropriate timeouts
             self.stream_client = httpx.AsyncClient(timeout=self.connection_timeout)
             self.send_client = httpx.AsyncClient(timeout=self.default_timeout)
             
-            # Connect to SSE stream with smart URL construction
+            # Connect to SSE stream
             sse_url = self._construct_sse_url(self.url)
-            logger.debug("Connecting to SSE: %s", sse_url)
+            logger.debug("Connecting to SSE endpoint: %s", sse_url)
             
             self.sse_stream_context = self.stream_client.stream(
                 'GET', sse_url, headers=self._get_headers()
@@ -116,28 +151,37 @@ class SSETransport(MCPBaseTransport):
             self.sse_response = await self.sse_stream_context.__aenter__()
             
             if self.sse_response.status_code != 200:
-                logger.error("SSE connection failed: %s", self.sse_response.status_code)
+                logger.error("SSE connection failed with status: %s", self.sse_response.status_code)
+                await self._cleanup()
                 return False
             
             logger.debug("SSE streaming connection established")
             
             # Start SSE processing task
-            self.sse_task = asyncio.create_task(self._process_sse_stream())
+            self.sse_task = asyncio.create_task(
+                self._process_sse_stream(), 
+                name="sse_stream_processor"
+            )
             
-            # Wait for session discovery
+            # Wait for session discovery with timeout
             logger.debug("Waiting for session discovery...")
-            for i in range(50):  # 5 seconds max
-                if self.message_url:
-                    break
+            session_timeout = 5.0  # 5 seconds max for session discovery
+            session_start = time.time()
+            
+            while not self.message_url and (time.time() - session_start) < session_timeout:
                 await asyncio.sleep(0.1)
             
             if not self.message_url:
-                logger.error("Failed to get session info from SSE")
+                logger.error("Failed to discover session endpoint within %.1fs", session_timeout)
+                await self._cleanup()
                 return False
             
-            logger.debug("Session ready: %s", self.session_id)
+            if self.enable_metrics:
+                self._metrics["session_discoveries"] += 1
+            
+            logger.debug("Session endpoint discovered: %s", self.session_id)
             
-            # Now do MCP initialization
+            # Perform MCP initialization handshake
             try:
                 init_response = await self._send_request("initialize", {
                     "protocolVersion": "2024-11-05",
@@ -149,18 +193,25 @@ class SSETransport(MCPBaseTransport):
                 })
                 
                 if 'error' in init_response:
-                    logger.error("Initialize failed: %s", init_response['error'])
+                    logger.error("MCP initialize failed: %s", init_response['error'])
+                    await self._cleanup()
                     return False
                 
                 # Send initialized notification
                 await self._send_notification("notifications/initialized")
                 
                 self._initialized = True
-                logger.debug("SSE transport initialized successfully")
+                
+                if self.enable_metrics:
+                    init_time = time.time() - start_time
+                    self._metrics["initialization_time"] = init_time
+                
+                logger.debug("SSE transport initialized successfully in %.3fs", time.time() - start_time)
                 return True
                 
             except Exception as e:
-                logger.error("MCP initialization failed: %s", e)
+                logger.error("MCP handshake failed: %s", e)
+                await self._cleanup()
                 return False
                 
         except Exception as e:
@@ -169,58 +220,94 @@ class SSETransport(MCPBaseTransport):
             return False
 
     async def _process_sse_stream(self):
-        """Process the persistent SSE stream."""
+        """
+        Process the persistent SSE stream for responses and session discovery.
+        
+        FIXED: Supports both old format (/messages/) and new event-based format
+        (event: endpoint + data: https://...) for session discovery.
+        """
         try:
             logger.debug("Starting SSE stream processing...")
             
+            current_event = None  # Track current event type
+            
             async for line in self.sse_response.aiter_lines():
                 line = line.strip()
                 if not line:
                     continue
                 
-                # Handle session endpoint discovery
-                if not self.message_url and line.startswith('data:') and '/messages/' in line:
-                    endpoint_path = line.split(':', 1)[1].strip()
-                    self.message_url = f"{self.url}{endpoint_path}"
+                # Handle event type declarations
+                if line.startswith('event:'):
+                    current_event = line.split(':', 1)[1].strip()
+                    logger.debug("SSE event type: %s", current_event)
+                    continue
+                
+                # Handle session endpoint discovery (BOTH FORMATS)
+                if not self.message_url and line.startswith('data:'):
+                    data_part = line.split(':', 1)[1].strip()
                     
-                    if 'session_id=' in endpoint_path:
-                        self.session_id = endpoint_path.split('session_id=')[1].split('&')[0]
+                    # NEW FORMAT: event: endpoint + data: https://...
+                    if current_event == "endpoint" and data_part.startswith('http'):
+                        self.message_url = data_part
+                        
+                        # Extract session ID from URL if present
+                        if 'session_id=' in data_part:
+                            self.session_id = data_part.split('session_id=')[1].split('&')[0]
+                        
+                        logger.debug("Session endpoint discovered via event format: %s", self.session_id)
+                        continue
                     
-                    logger.debug("Got session info: %s", self.session_id)
-                    continue
+                    # OLD FORMAT: data: /messages/... (backwards compatibility)
+                    elif '/messages/' in data_part:
+                        endpoint_path = data_part
+                        self.message_url = f"{self.url}{endpoint_path}"
+                        
+                        # Extract session ID if present
+                        if 'session_id=' in endpoint_path:
+                            self.session_id = endpoint_path.split('session_id=')[1].split('&')[0]
+                        
+                        logger.debug("Session endpoint discovered via old format: %s", self.session_id)
+                        continue
                 
                 # Handle JSON-RPC responses
                 if line.startswith('data:'):
                     data_part = line.split(':', 1)[1].strip()
                     
-                    # Skip pings and empty data
-                    if not data_part or data_part.startswith('ping'):
+                    # Skip keepalive pings and empty data
+                    if not data_part or data_part.startswith('ping') or data_part in ('{}', '[]'):
                         continue
                     
                     try:
                         response_data = json.loads(data_part)
                         
+                        # Handle JSON-RPC responses with request IDs
                         if 'jsonrpc' in response_data and 'id' in response_data:
                             request_id = str(response_data['id'])
                             
-                            # Resolve pending request
+                            # Resolve pending request if found
                             if request_id in self.pending_requests:
                                 future = self.pending_requests.pop(request_id)
                                 if not future.done():
                                     future.set_result(response_data)
-                                    logger.debug("Resolved request: %s", request_id)
+                                    logger.debug("Resolved request ID: %s", request_id)
                     
-                    except json.JSONDecodeError:
-                        pass  # Not JSON, ignore
+                    except json.JSONDecodeError as e:
+                        logger.debug("Non-JSON data in SSE stream (ignoring): %s", e)
+                
+                # Reset event type after processing data (only if we processed JSON-RPC)
+                if line.startswith('data:') and current_event not in ("endpoint",):
+                    current_event = None
                         
         except Exception as e:
-            logger.error("SSE stream error: %s", e)
+            if self.enable_metrics:
+                self._metrics["stream_errors"] += 1
+            logger.error("SSE stream processing error: %s", e)
 
     async def _send_request(self, method: str, params: Dict[str, Any] = None, 
                            timeout: Optional[float] = None) -> Dict[str, Any]:
-        """Send request and wait for async response."""
+        """Send JSON-RPC request and wait for async response via SSE."""
         if not self.message_url:
-            raise RuntimeError("Not connected")
+            raise RuntimeError("SSE transport not connected - no message URL")
         
         request_id = str(uuid.uuid4())
         message = {
@@ -230,12 +317,12 @@ class SSETransport(MCPBaseTransport):
             "params": params or {}
         }
         
-        # Create future for response
+        # Create future for async response
         future = asyncio.Future()
         self.pending_requests[request_id] = future
         
         try:
-            # Send message
+            # Send HTTP POST request
             headers = {
                 'Content-Type': 'application/json',
                 **self._get_headers()
@@ -248,9 +335,9 @@ class SSETransport(MCPBaseTransport):
             )
             
             if response.status_code == 202:
-                # Wait for async response
-                timeout = timeout or self.default_timeout
-                result = await asyncio.wait_for(future, timeout=timeout)
+                # Async response - wait for result via SSE
+                request_timeout = timeout or self.default_timeout
+                result = await asyncio.wait_for(future, timeout=request_timeout)
                 return result
             elif response.status_code == 200:
                 # Immediate response
@@ -258,7 +345,7 @@ class SSETransport(MCPBaseTransport):
                 return response.json()
             else:
                 self.pending_requests.pop(request_id, None)
-                raise RuntimeError(f"Request failed: {response.status_code}")
+                raise RuntimeError(f"HTTP request failed with status: {response.status_code}")
                 
         except asyncio.TimeoutError:
             self.pending_requests.pop(request_id, None)
@@ -268,9 +355,9 @@ class SSETransport(MCPBaseTransport):
             raise
 
     async def _send_notification(self, method: str, params: Dict[str, Any] = None):
-        """Send notification (no response expected)."""
+        """Send JSON-RPC notification (no response expected)."""
         if not self.message_url:
-            raise RuntimeError("Not connected")
+            raise RuntimeError("SSE transport not connected - no message URL")
         
         message = {
             "jsonrpc": "2.0",
@@ -283,30 +370,46 @@ class SSETransport(MCPBaseTransport):
             **self._get_headers()
         }
         
-        await self.send_client.post(
+        response = await self.send_client.post(
             self.message_url,
             headers=headers,
             json=message
         )
+        
+        if response.status_code not in (200, 202):
+            logger.warning("Notification failed with status: %s", response.status_code)
 
     async def send_ping(self) -> bool:
-        """Send ping to check connection."""
+        """Send ping to check connection health."""
         if not self._initialized:
             return False
         
+        start_time = time.time()
         try:
-            # Your server might not support ping, so we'll just check if we can list tools
+            # Use tools/list as a lightweight ping since not all servers support ping
             response = await self._send_request("tools/list", {}, timeout=5.0)
+            
+            if self.enable_metrics:
+                ping_time = time.time() - start_time
+                self._metrics["last_ping_time"] = ping_time
+                logger.debug("SSE ping completed in %.3fs", ping_time)
+            
             return 'error' not in response
-        except Exception:
+        except Exception as e:
+            logger.debug("SSE ping failed: %s", e)
             return False
 
+    def is_connected(self) -> bool:
+        """Check if the transport is connected and ready."""
+        return self._initialized and self.session_id is not None
+
     async def get_tools(self) -> List[Dict[str, Any]]:
-        """Get tools list."""
+        """Get list of available tools from the server."""
         if not self._initialized:
             logger.error("Cannot get tools: transport not initialized")
             return []
         
+        start_time = time.time()
         try:
             response = await self._send_request("tools/list", {})
             
@@ -315,7 +418,11 @@ class SSETransport(MCPBaseTransport):
                 return []
             
             tools = response.get('result', {}).get('tools', [])
-            logger.debug("Retrieved %d tools", len(tools))
+            
+            if self.enable_metrics:
+                response_time = time.time() - start_time
+                logger.debug("Retrieved %d tools in %.3fs", len(tools), response_time)
+            
             return tools
             
         except Exception as e:
@@ -324,15 +431,19 @@ class SSETransport(MCPBaseTransport):
 
     async def call_tool(self, tool_name: str, arguments: Dict[str, Any], 
                        timeout: Optional[float] = None) -> Dict[str, Any]:
-        """Call a tool."""
+        """Execute a tool with the given arguments."""
         if not self._initialized:
             return {
                 "isError": True,
                 "error": "Transport not initialized"
             }
 
+        start_time = time.time()
+        if self.enable_metrics:
+            self._metrics["total_calls"] += 1
+
         try:
-            logger.debug("Calling tool %s with args: %s", tool_name, arguments)
+            logger.debug("Calling tool '%s' with arguments: %s", tool_name, arguments)
             
             response = await self._send_request(
                 "tools/call",
@@ -344,58 +455,56 @@ class SSETransport(MCPBaseTransport):
             )
             
             if 'error' in response:
+                if self.enable_metrics:
+                    self._update_metrics(time.time() - start_time, False)
+                
                 return {
                     "isError": True,
                     "error": response['error'].get('message', 'Unknown error')
                 }
             
-            # Extract result
+            # Extract and normalize result using base class method
             result = response.get('result', {})
+            normalized_result = self._normalize_mcp_response({"result": result})
             
-            # Handle content format
-            if 'content' in result:
-                content = result['content']
-                if isinstance(content, list) and len(content) == 1:
-                    content_item = content[0]
-                    if isinstance(content_item, dict) and content_item.get('type') == 'text':
-                        text_content = content_item.get('text', '')
-                        try:
-                            # Try to parse as JSON
-                            parsed_content = json.loads(text_content)
-                            return {
-                                "isError": False,
-                                "content": parsed_content
-                            }
-                        except json.JSONDecodeError:
-                            return {
-                                "isError": False,
-                                "content": text_content
-                            }
-                
-                return {
-                    "isError": False,
-                    "content": content
-                }
+            if self.enable_metrics:
+                self._update_metrics(time.time() - start_time, True)
             
-            return {
-                "isError": False,
-                "content": result
-            }
+            return normalized_result
             
         except asyncio.TimeoutError:
+            if self.enable_metrics:
+                self._update_metrics(time.time() - start_time, False)
+            
             return {
                 "isError": True,
                 "error": "Tool execution timed out"
             }
         except Exception as e:
-            logger.error("Error calling tool %s: %s", tool_name, e)
+            if self.enable_metrics:
+                self._update_metrics(time.time() - start_time, False)
+            
+            logger.error("Error calling tool '%s': %s", tool_name, e)
             return {
                 "isError": True,
                 "error": str(e)
             }
 
+    def _update_metrics(self, response_time: float, success: bool) -> None:
+        """Update performance metrics."""
+        if success:
+            self._metrics["successful_calls"] += 1
+        else:
+            self._metrics["failed_calls"] += 1
+            
+        self._metrics["total_time"] += response_time
+        if self._metrics["total_calls"] > 0:
+            self._metrics["avg_response_time"] = (
+                self._metrics["total_time"] / self._metrics["total_calls"]
+            )
+
     async def list_resources(self) -> Dict[str, Any]:
-        """List resources."""
+        """List available resources from the server."""
         if not self._initialized:
             return {}
         
@@ -405,11 +514,12 @@ class SSETransport(MCPBaseTransport):
                 logger.debug("Resources not supported: %s", response['error'])
                 return {}
             return response.get('result', {})
-        except Exception:
+        except Exception as e:
+            logger.debug("Error listing resources: %s", e)
             return {}
 
     async def list_prompts(self) -> Dict[str, Any]:
-        """List prompts."""
+        """List available prompts from the server."""
         if not self._initialized:
             return {}
         
@@ -419,59 +529,104 @@ class SSETransport(MCPBaseTransport):
                 logger.debug("Prompts not supported: %s", response['error'])
                 return {}
             return response.get('result', {})
-        except Exception:
+        except Exception as e:
+            logger.debug("Error listing prompts: %s", e)
             return {}
 
     async def close(self) -> None:
-        """Close the transport."""
+        """Close the transport and clean up resources."""
+        if not self._initialized:
+            return
+        
+        # Log final metrics
+        if self.enable_metrics and self._metrics["total_calls"] > 0:
+            logger.debug(
+                "SSE transport closing - Total calls: %d, Success rate: %.1f%%, Avg response time: %.3fs",
+                self._metrics["total_calls"],
+                (self._metrics["successful_calls"] / self._metrics["total_calls"] * 100),
+                self._metrics["avg_response_time"]
+            )
+        
         await self._cleanup()
 
     async def _cleanup(self) -> None:
-        """Clean up resources."""
-        if self.sse_task:
+        """Clean up all resources and reset state."""
+        # Cancel SSE processing task
+        if self.sse_task and not self.sse_task.done():
             self.sse_task.cancel()
             try:
                 await self.sse_task
             except asyncio.CancelledError:
                 pass
         
+        # Close SSE stream context
         if self.sse_stream_context:
             try:
                 await self.sse_stream_context.__aexit__(None, None, None)
-            except Exception:
-                pass
+            except Exception as e:
+                logger.debug("Error closing SSE stream: %s", e)
         
+        # Close HTTP clients
         if self.stream_client:
             await self.stream_client.aclose()
         
         if self.send_client:
             await self.send_client.aclose()
         
+        # Cancel any pending requests
+        for request_id, future in self.pending_requests.items():
+            if not future.done():
+                future.cancel()
+        
+        # Reset state
         self._initialized = False
         self.session_id = None
         self.message_url = None
         self.pending_requests.clear()
+        self.sse_task = None
+        self.sse_response = None
+        self.sse_stream_context = None
+        self.stream_client = None
+        self.send_client = None
+
+    # ------------------------------------------------------------------ #
+    #  Metrics and monitoring (consistent with other transports)        #
+    # ------------------------------------------------------------------ #
+    def get_metrics(self) -> Dict[str, Any]:
+        """Get performance and connection metrics."""
+        return self._metrics.copy()
+
+    def reset_metrics(self) -> None:
+        """Reset performance metrics."""
+        self._metrics = {
+            "total_calls": 0,
+            "successful_calls": 0,
+            "failed_calls": 0,
+            "total_time": 0.0,
+            "avg_response_time": 0.0,
+            "last_ping_time": self._metrics.get("last_ping_time"),
+            "initialization_time": self._metrics.get("initialization_time"),
+            "session_discoveries": self._metrics.get("session_discoveries", 0),
+            "stream_errors": 0
+        }
 
+    # ------------------------------------------------------------------ #
+    #  Backward compatibility                                            #
+    # ------------------------------------------------------------------ #
     def get_streams(self) -> List[tuple]:
-        """Not applicable for this transport."""
+        """SSE transport doesn't expose raw streams."""
         return []
 
-    def is_connected(self) -> bool:
-        """Check if connected."""
-        return self._initialized and self.session_id is not None
-
+    # ------------------------------------------------------------------ #
+    #  Context manager support                                           #
+    # ------------------------------------------------------------------ #
     async def __aenter__(self):
-        """Context manager support."""
+        """Context manager entry."""
         success = await self.initialize()
         if not success:
-            raise RuntimeError("Failed to initialize SSE transport")
+            raise RuntimeError("Failed to initialize SSETransport")
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         """Context manager cleanup."""
-        await self.close()
-
-    def __repr__(self) -> str:
-        """String representation."""
-        status = "initialized" if self._initialized else "not initialized"
-        return f"SSETransport(status={status}, url={self.url}, session={self.session_id})"
+        await self.close()