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

chuk_tool_processor/mcp/transport/stdio_transport.py

@@ -1,9 +1,11 @@
-# chuk_tool_processor/mcp/transport/stdio_transport.py
+# chuk_tool_processor/mcp/transport/stdio_transport.py - ENHANCED
 from __future__ import annotations
 
 import asyncio
 import json
 import time
+import signal
+import psutil
 from typing import Dict, Any, List, Optional
 import logging
 
@@ -23,22 +25,24 @@ class StdioTransport(MCPBaseTransport):
     """
     STDIO transport for MCP communication using process pipes.
     
-    This transport uses subprocess communication via stdin/stdout pipes
-    to communicate with MCP servers.
+    ENHANCED: Now matches SSE transport robustness with improved process 
+    management, health monitoring, and comprehensive error handling.
     """
 
     def __init__(self, server_params, 
                  connection_timeout: float = 30.0,
                  default_timeout: float = 30.0,
-                 enable_metrics: bool = True):
+                 enable_metrics: bool = True,
+                 process_monitor: bool = True):  # NEW
         """
-        Initialize STDIO transport.
+        Initialize STDIO transport with enhanced configuration.
         
         Args:
             server_params: Server parameters (dict or StdioParameters object)
             connection_timeout: Timeout for initial connection setup
             default_timeout: Default timeout for operations
             enable_metrics: Whether to track performance metrics
+            process_monitor: Whether to monitor subprocess health (NEW)
         """
         # Convert dict to StdioParameters if needed
         if isinstance(server_params, dict):
@@ -53,13 +57,21 @@ class StdioTransport(MCPBaseTransport):
         self.connection_timeout = connection_timeout
         self.default_timeout = default_timeout
         self.enable_metrics = enable_metrics
+        self.process_monitor = process_monitor  # NEW
         
         # Connection state
         self._context = None
         self._streams = None
         self._initialized = False
         
-        # Performance metrics (consistent with other transports)
+        # Process monitoring (NEW - like SSE's health monitoring)
+        self._process_id = None
+        self._process_start_time = None
+        self._last_successful_ping = None
+        self._consecutive_failures = 0
+        self._max_consecutive_failures = 3
+        
+        # Enhanced performance metrics (like SSE)
         self._metrics = {
             "total_calls": 0,
             "successful_calls": 0,
@@ -69,14 +81,76 @@ class StdioTransport(MCPBaseTransport):
             "last_ping_time": None,
             "initialization_time": None,
             "process_restarts": 0,
-            "pipe_errors": 0
+            "pipe_errors": 0,
+            "process_crashes": 0,  # NEW
+            "recovery_attempts": 0,  # NEW
+            "memory_usage_mb": 0.0,  # NEW
+            "cpu_percent": 0.0,  # NEW
         }
         
         logger.debug("STDIO transport initialized for command: %s", 
                     getattr(self.server_params, 'command', 'unknown'))
 
+    async def _get_process_info(self) -> Optional[Dict[str, Any]]:
+        """Get process information for monitoring (NEW)."""
+        if not self._process_id or not self.process_monitor:
+            return None
+        
+        try:
+            # FIXED: Validate PID is a real integer before using psutil
+            if not isinstance(self._process_id, int) or self._process_id <= 0:
+                return None
+                
+            process = psutil.Process(self._process_id)
+            if process.is_running():
+                memory_info = process.memory_info()
+                return {
+                    "pid": self._process_id,
+                    "status": process.status(),
+                    "memory_mb": memory_info.rss / 1024 / 1024,
+                    "cpu_percent": process.cpu_percent(),
+                    "create_time": process.create_time(),
+                    "uptime": time.time() - self._process_start_time if self._process_start_time else 0
+                }
+        except (psutil.NoSuchProcess, psutil.AccessDenied, AttributeError, TypeError, ValueError):
+            # FIXED: Handle all possible errors including TypeError from mock objects
+            pass
+        return None
+
+    async def _monitor_process_health(self) -> bool:
+        """Monitor subprocess health (NEW - like SSE's health monitoring)."""
+        if not self.process_monitor:
+            return True
+        
+        # FIXED: Check if process_id is valid before monitoring
+        if not self._process_id or not isinstance(self._process_id, int) or self._process_id <= 0:
+            return True  # No monitoring if no valid PID
+        
+        process_info = await self._get_process_info()
+        if not process_info:
+            logger.debug("Process monitoring unavailable (may be in test environment)")
+            return True  # Don't fail in test environments
+        
+        # Update metrics with process info
+        if self.enable_metrics:
+            self._metrics["memory_usage_mb"] = process_info["memory_mb"]
+            self._metrics["cpu_percent"] = process_info["cpu_percent"]
+        
+        # Check for concerning process states
+        status = process_info.get("status", "unknown")
+        if status in ["zombie", "dead"]:
+            logger.error("Process is in %s state", status)
+            return False
+        
+        # Check for excessive memory usage (warn at 1GB)
+        memory_mb = process_info.get("memory_mb", 0)
+        if memory_mb > 1024:
+            logger.warning("Process using excessive memory: %.1f MB", memory_mb)
+        
+        return True
+
     async def initialize(self) -> bool:
-        """Initialize by delegating to chuk-mcp with timeout protection."""
+        """Enhanced initialization with process monitoring."""
         if self._initialized:
             logger.warning("Transport already initialized")
             return True
@@ -93,6 +167,12 @@ class StdioTransport(MCPBaseTransport):
                 timeout=self.connection_timeout
             )
             
+            # Capture process information for monitoring (NEW)
+            if self.process_monitor and hasattr(self._context, '_process'):
+                self._process_id = getattr(self._context._process, 'pid', None)
+                self._process_start_time = time.time()
+                logger.debug("Subprocess PID: %s", self._process_id)
+            
             # Send initialize message with timeout
             init_result = await asyncio.wait_for(
                 send_initialize(*self._streams),
@@ -100,14 +180,36 @@ class StdioTransport(MCPBaseTransport):
             )
             
             if init_result:
-                self._initialized = True
+                # Enhanced health verification (like SSE)
+                logger.debug("Verifying connection with ping...")
+                ping_start = time.time()
+                ping_success = await asyncio.wait_for(
+                    send_ping(*self._streams),
+                    timeout=10.0
+                )
+                ping_time = time.time() - ping_start
                 
-                if self.enable_metrics:
-                    init_time = time.time() - start_time
-                    self._metrics["initialization_time"] = init_time
-                
-                logger.debug("STDIO transport initialized successfully in %.3fs", time.time() - start_time)
-                return True
+                if ping_success:
+                    self._initialized = True
+                    self._last_successful_ping = time.time()
+                    self._consecutive_failures = 0
+                    
+                    if self.enable_metrics:
+                        init_time = time.time() - start_time
+                        self._metrics["initialization_time"] = init_time
+                        self._metrics["last_ping_time"] = ping_time
+                    
+                    logger.debug("STDIO transport initialized successfully in %.3fs (ping: %.3fs)", 
+                               time.time() - start_time, ping_time)
+                    return True
+                else:
+                    logger.warning("STDIO connection established but ping failed")
+                    # Still consider it initialized
+                    self._initialized = True
+                    self._consecutive_failures = 1
+                    if self.enable_metrics:
+                        self._metrics["initialization_time"] = time.time() - start_time
+                    return True
             else:
                 logger.error("STDIO initialization failed")
                 await self._cleanup()
@@ -116,24 +218,54 @@ class StdioTransport(MCPBaseTransport):
         except asyncio.TimeoutError:
             logger.error("STDIO initialization timed out after %ss", self.connection_timeout)
             await self._cleanup()
+            if self.enable_metrics:
+                self._metrics["process_crashes"] += 1
             return False
         except Exception as e:
             logger.error("Error initializing STDIO transport: %s", e)
             await self._cleanup()
+            if self.enable_metrics:
+                self._metrics["process_crashes"] += 1
+            return False
+
+    async def _attempt_recovery(self) -> bool:
+        """Attempt to recover from process/connection issues (NEW)."""
+        if self.enable_metrics:
+            self._metrics["recovery_attempts"] += 1
+            self._metrics["process_restarts"] += 1
+        
+        logger.warning("Attempting STDIO process recovery...")
+        
+        try:
+            # Force cleanup of existing process
+            await self._cleanup()
+            
+            # Brief delay before restart
+            await asyncio.sleep(1.0)
+            
+            # Re-initialize
+            return await self.initialize()
+        except Exception as e:
+            logger.error("Recovery attempt failed: %s", e)
             return False
 
     async def close(self) -> None:
-        """Close by delegating to chuk-mcp context manager with enhanced cleanup."""
+        """Enhanced close with process monitoring and metrics."""
         if not self._initialized:
             return
         
-        # Log final metrics
+        # Enhanced metrics logging (like SSE)
         if self.enable_metrics and self._metrics["total_calls"] > 0:
+            success_rate = (self._metrics["successful_calls"] / self._metrics["total_calls"] * 100)
             logger.debug(
-                "STDIO transport closing - Total calls: %d, Success rate: %.1f%%, Avg response time: %.3fs",
+                "STDIO transport closing - Calls: %d, Success: %.1f%%, "
+                "Avg time: %.3fs, Restarts: %d, Crashes: %d, Memory: %.1f MB",
                 self._metrics["total_calls"],
-                (self._metrics["successful_calls"] / self._metrics["total_calls"] * 100),
-                self._metrics["avg_response_time"]
+                success_rate,
+                self._metrics["avg_response_time"],
+                self._metrics["process_restarts"],
+                self._metrics["process_crashes"],
+                self._metrics["memory_usage_mb"]
             )
         
         if self._context:
@@ -146,16 +278,44 @@ class StdioTransport(MCPBaseTransport):
                 await self._cleanup()
 
     async def _cleanup(self) -> None:
-        """Clean up internal state."""
+        """Enhanced cleanup with process termination."""
+        # Attempt graceful process termination if we have a PID
+        if self._process_id and self.process_monitor:
+            try:
+                # FIXED: Validate PID is a real integer before using psutil
+                if isinstance(self._process_id, int) and self._process_id > 0:
+                    process = psutil.Process(self._process_id)
+                    if process.is_running():
+                        logger.debug("Terminating subprocess %s", self._process_id)
+                        process.terminate()
+                        
+                        # Wait briefly for graceful termination
+                        try:
+                            process.wait(timeout=2.0)
+                        except psutil.TimeoutExpired:
+                            logger.warning("Process did not terminate gracefully, killing...")
+                            process.kill()
+            except (psutil.NoSuchProcess, psutil.AccessDenied, TypeError, ValueError):
+                # FIXED: Handle all possible errors including TypeError from mock objects
+                logger.debug("Could not terminate process %s (may be mock or already dead)", self._process_id)
+        
         self._context = None
         self._streams = None
         self._initialized = False
+        self._process_id = None
+        self._process_start_time = None
 
     async def send_ping(self) -> bool:
-        """Send ping with performance tracking."""
+        """Enhanced ping with process health monitoring."""
         if not self._initialized:
             return False
         
+        # Check process health first (NEW) - but only if we have a real process
+        if self.process_monitor and self._process_id and isinstance(self._process_id, int):
+            if not await self._monitor_process_health():
+                self._consecutive_failures += 1
+                return False
+        
         start_time = time.time()
         try:
             result = await asyncio.wait_for(
@@ -163,27 +323,45 @@ class StdioTransport(MCPBaseTransport):
                 timeout=self.default_timeout
             )
             
+            success = bool(result)
+            
+            if success:
+                self._last_successful_ping = time.time()
+                self._consecutive_failures = 0
+            else:
+                self._consecutive_failures += 1
+            
             if self.enable_metrics:
                 ping_time = time.time() - start_time
                 self._metrics["last_ping_time"] = ping_time
-                logger.debug("STDIO ping completed in %.3fs: %s", ping_time, result)
+                logger.debug("STDIO ping completed in %.3fs: %s", ping_time, success)
             
-            return bool(result)
+            return success
         except asyncio.TimeoutError:
             logger.error("STDIO ping timed out")
+            self._consecutive_failures += 1
             return False
         except Exception as e:
             logger.error("STDIO ping failed: %s", e)
+            self._consecutive_failures += 1
             if self.enable_metrics:
                 self._metrics["pipe_errors"] += 1
             return False
 
     def is_connected(self) -> bool:
-        """Check connection status."""
-        return self._initialized and self._streams is not None
+        """Enhanced connection status check (like SSE)."""
+        if not self._initialized or not self._streams:
+            return False
+        
+        # Check for too many consecutive failures (like SSE)
+        if self._consecutive_failures >= self._max_consecutive_failures:
+            logger.warning("Connection marked unhealthy after %d failures", self._consecutive_failures)
+            return False
+        
+        return True
 
     async def get_tools(self) -> List[Dict[str, Any]]:
-        """Get tools list with performance tracking."""
+        """Enhanced tools retrieval with recovery."""
         if not self._initialized:
             logger.error("Cannot get tools: transport not initialized")
             return []
@@ -204,6 +382,9 @@ class StdioTransport(MCPBaseTransport):
                 logger.warning("Unexpected tools response type: %s", type(response))
                 tools = []
             
+            # Reset failure count on success
+            self._consecutive_failures = 0
+            
             if self.enable_metrics:
                 response_time = time.time() - start_time
                 logger.debug("Retrieved %d tools in %.3fs", len(tools), response_time)
@@ -212,16 +393,18 @@ class StdioTransport(MCPBaseTransport):
             
         except asyncio.TimeoutError:
             logger.error("Get tools timed out")
+            self._consecutive_failures += 1
             return []
         except Exception as e:
             logger.error("Error getting tools: %s", e)
+            self._consecutive_failures += 1
             if self.enable_metrics:
                 self._metrics["pipe_errors"] += 1
             return []
 
     async def call_tool(self, tool_name: str, arguments: Dict[str, Any], 
                        timeout: Optional[float] = None) -> Dict[str, Any]:
-        """Call tool with timeout support and performance tracking."""
+        """Enhanced tool calling with recovery and process monitoring."""
         if not self._initialized:
             return {"isError": True, "error": "Transport not initialized"}
 
@@ -229,11 +412,22 @@ class StdioTransport(MCPBaseTransport):
         start_time = time.time()
         
         if self.enable_metrics:
-            self._metrics["total_calls"] += 1  # FIXED: INCREMENT FIRST
+            self._metrics["total_calls"] += 1
 
         try:
             logger.debug("Calling tool '%s' with timeout %ss", tool_name, tool_timeout)
             
+            # Enhanced connection check with recovery attempt
+            if not self.is_connected():
+                logger.warning("Connection unhealthy, attempting recovery...")
+                if not await self._attempt_recovery():
+                    if self.enable_metrics:
+                        self._update_metrics(time.time() - start_time, False)
+                    return {
+                        "isError": True,
+                        "error": "Failed to recover connection"
+                    }
+            
             response = await asyncio.wait_for(
                 send_tools_call(*self._streams, tool_name, arguments),
                 timeout=tool_timeout
@@ -242,18 +436,24 @@ class StdioTransport(MCPBaseTransport):
             response_time = time.time() - start_time
             result = self._normalize_mcp_response(response)
             
+            # Reset failure count and update health on success
+            self._consecutive_failures = 0
+            self._last_successful_ping = time.time()
+            
             if self.enable_metrics:
                 self._update_metrics(response_time, not result.get("isError", False))
             
             if not result.get("isError", False):
                 logger.debug("Tool '%s' completed successfully in %.3fs", tool_name, response_time)
             else:
-                logger.warning("Tool '%s' failed in %.3fs: %s", tool_name, response_time, result.get('error', 'Unknown error'))
+                logger.warning("Tool '%s' failed in %.3fs: %s", tool_name, response_time, 
+                             result.get('error', 'Unknown error'))
             
             return result
             
         except asyncio.TimeoutError:
             response_time = time.time() - start_time
+            self._consecutive_failures += 1
             if self.enable_metrics:
                 self._update_metrics(response_time, False)
             
@@ -265,10 +465,20 @@ class StdioTransport(MCPBaseTransport):
             }
         except Exception as e:
             response_time = time.time() - start_time
+            self._consecutive_failures += 1
             if self.enable_metrics:
                 self._update_metrics(response_time, False)
                 self._metrics["pipe_errors"] += 1
             
+            # Enhanced process error detection
+            error_str = str(e).lower()
+            if any(indicator in error_str for indicator in 
+                   ["broken pipe", "process", "eof", "connection", "died"]):
+                logger.warning("Process error detected: %s", e)
+                self._initialized = False
+                if self.enable_metrics:
+                    self._metrics["process_crashes"] += 1
+            
             error_msg = f"Tool execution failed: {str(e)}"
             logger.error("Tool '%s' error: %s", tool_name, error_msg)
             return {
@@ -277,14 +487,13 @@ class StdioTransport(MCPBaseTransport):
             }
 
     def _update_metrics(self, response_time: float, success: bool) -> None:
-        """Update performance metrics."""
+        """Enhanced metrics tracking (like SSE)."""
         if success:
             self._metrics["successful_calls"] += 1
         else:
             self._metrics["failed_calls"] += 1
             
         self._metrics["total_time"] += response_time
-        # FIXED: Only calculate average if we have total calls
         if self._metrics["total_calls"] > 0:
             self._metrics["avg_response_time"] = (
                 self._metrics["total_time"] / self._metrics["total_calls"]
@@ -292,10 +501,10 @@ class StdioTransport(MCPBaseTransport):
 
     def _normalize_mcp_response(self, response: Dict[str, Any]) -> Dict[str, Any]:
         """
-        Normalize response using shared base class logic with STDIO-specific handling.
+        Enhanced response normalization with STDIO-specific handling.
         
-        STDIO has special requirements for preserving string representations
-        of numeric values for backward compatibility.
+        STDIO preserves string representations of numeric values for 
+        backward compatibility with existing tests.
         """
         # Handle explicit error in response
         if "error" in response:
@@ -318,7 +527,7 @@ class StdioTransport(MCPBaseTransport):
 
     def _extract_stdio_content(self, content_list: Any) -> Any:
         """
-        Extract content with STDIO-specific string preservation logic.
+        Enhanced content extraction with STDIO-specific string preservation.
         
         STDIO transport preserves string representations of numeric values
         for backward compatibility with existing tests.
@@ -348,7 +557,7 @@ class StdioTransport(MCPBaseTransport):
         return content_list
 
     async def list_resources(self) -> Dict[str, Any]:
-        """List resources with error handling."""
+        """Enhanced resource listing with error handling."""
         if not self._initialized:
             return {}
         try:
@@ -356,16 +565,19 @@ class StdioTransport(MCPBaseTransport):
                 send_resources_list(*self._streams),
                 timeout=self.default_timeout
             )
+            self._consecutive_failures = 0  # Reset on success
             return response if isinstance(response, dict) else {}
         except asyncio.TimeoutError:
             logger.error("List resources timed out")
+            self._consecutive_failures += 1
             return {}
         except Exception as e:
             logger.debug("Error listing resources: %s", e)
+            self._consecutive_failures += 1
             return {}
 
     async def list_prompts(self) -> Dict[str, Any]:
-        """List prompts with error handling."""
+        """Enhanced prompt listing with error handling."""
         if not self._initialized:
             return {}
         try:
@@ -373,12 +585,15 @@ class StdioTransport(MCPBaseTransport):
                 send_prompts_list(*self._streams),
                 timeout=self.default_timeout
             )
+            self._consecutive_failures = 0  # Reset on success
             return response if isinstance(response, dict) else {}
         except asyncio.TimeoutError:
             logger.error("List prompts timed out")
+            self._consecutive_failures += 1
             return {}
         except Exception as e:
             logger.debug("Error listing prompts: %s", e)
+            self._consecutive_failures += 1
             return {}
 
     async def read_resource(self, uri: str) -> Dict[str, Any]:
@@ -390,9 +605,15 @@ class StdioTransport(MCPBaseTransport):
                 send_resources_read(*self._streams, uri),
                 timeout=self.default_timeout
             )
+            self._consecutive_failures = 0  # Reset on success
             return response if isinstance(response, dict) else {}
+        except asyncio.TimeoutError:
+            logger.error("Read resource timed out")
+            self._consecutive_failures += 1
+            return {}
         except Exception as e:
             logger.debug("Error reading resource: %s", e)
+            self._consecutive_failures += 1
             return {}
 
     async def get_prompt(self, name: str, arguments: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
@@ -404,49 +625,63 @@ class StdioTransport(MCPBaseTransport):
                 send_prompts_get(*self._streams, name, arguments or {}),
                 timeout=self.default_timeout
             )
+            self._consecutive_failures = 0  # Reset on success
             return response if isinstance(response, dict) else {}
+        except asyncio.TimeoutError:
+            logger.error("Get prompt timed out")
+            self._consecutive_failures += 1
+            return {}
         except Exception as e:
             logger.debug("Error getting prompt: %s", e)
+            self._consecutive_failures += 1
             return {}
 
-    # ------------------------------------------------------------------ #
-    #  Metrics and monitoring (now consistent with other transports)    #
-    # ------------------------------------------------------------------ #
     def get_metrics(self) -> Dict[str, Any]:
-        """Get performance and connection metrics."""
-        return self._metrics.copy()
+        """Enhanced metrics with process and health information."""
+        metrics = self._metrics.copy()
+        metrics.update({
+            "is_connected": self.is_connected(),
+            "consecutive_failures": self._consecutive_failures,
+            "last_successful_ping": self._last_successful_ping,
+            "max_consecutive_failures": self._max_consecutive_failures,
+            "process_id": self._process_id,
+            "process_uptime": (time.time() - self._process_start_time) if self._process_start_time else 0,
+        })
+        return metrics
 
     def reset_metrics(self) -> None:
-        """Reset performance metrics."""
+        """Enhanced metrics reset preserving health and process state."""
+        preserved_init_time = self._metrics.get("initialization_time")
+        preserved_last_ping = self._metrics.get("last_ping_time")
+        preserved_restarts = self._metrics.get("process_restarts", 0)
+        
         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"),
-            "process_restarts": self._metrics.get("process_restarts", 0),
-            "pipe_errors": 0
+            "last_ping_time": preserved_last_ping,
+            "initialization_time": preserved_init_time,
+            "process_restarts": preserved_restarts,
+            "pipe_errors": 0,
+            "process_crashes": 0,
+            "recovery_attempts": 0,
+            "memory_usage_mb": 0.0,
+            "cpu_percent": 0.0,
         }
 
-    # ------------------------------------------------------------------ #
-    #  Backward compatibility                                            #
-    # ------------------------------------------------------------------ #
     def get_streams(self) -> List[tuple]:
-        """Provide streams for backward compatibility."""
+        """Enhanced streams access with connection check."""
         return [self._streams] if self._streams else []
 
-    # ------------------------------------------------------------------ #
-    #  Context manager support (now uses base class with fixed error)   #
-    # ------------------------------------------------------------------ #
     async def __aenter__(self):
-        """Context manager support."""
+        """Enhanced context manager entry."""
         success = await self.initialize()
         if not success:
-            raise RuntimeError("Failed to initialize StdioTransport")  # FIXED: message
+            raise RuntimeError("Failed to initialize StdioTransport")
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        """Context manager cleanup."""
+        """Enhanced context manager cleanup."""
         await self.close()

{chuk_tool_processor-0.6.9.dist-info → chuk_tool_processor-0.6.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chuk-tool-processor
-Version: 0.6.9
+Version: 0.6.11
 Summary: Async-native framework for registering, discovering, and executing tools referenced in LLM responses
 Author-email: CHUK Team <chrishayuk@somejunkmailbox.com>
 Maintainer-email: CHUK Team <chrishayuk@somejunkmailbox.com>
@@ -22,6 +22,7 @@ Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 Requires-Dist: chuk-mcp>=0.5.2
 Requires-Dist: dotenv>=0.9.9
+Requires-Dist: psutil>=7.0.0
 Requires-Dist: pydantic>=2.11.3
 Requires-Dist: uuid>=1.30