chuk-tool-processor 0.4.1__py3-none-any.whl → 0.5__py3-none-any.whl

chuk_tool_processor/mcp/mcp_tool.py

@@ -3,6 +3,8 @@
 """
 MCP tool shim that delegates execution to a StreamManager,
 handling its own lazy bootstrap when needed.
+
+FIXED: Added subprocess serialization support by implementing __getstate__ and __setstate__
 """
 from __future__ import annotations
 
@@ -24,12 +26,14 @@ class MCPTool:
 
     If no ``StreamManager`` is supplied the class will start one on first
     use via ``setup_mcp_stdio``.
+    
+    FIXED: Added serialization support for subprocess execution.
     """
 
     # ------------------------------------------------------------------ #
     def __init__(
         self,
-        tool_name: str,
+        tool_name: str = "",
         stream_manager: Optional[StreamManager] = None,
         *,
         cfg_file: str = "",
@@ -38,6 +42,13 @@ class MCPTool:
         namespace: str = "stdio",
         default_timeout: Optional[float] = None
     ) -> None:
+        if not tool_name:
+            raise ValueError(
+                "MCPTool requires a tool_name. "
+                "This error usually occurs during subprocess serialization. "
+                "Make sure the tool is properly registered with a name."
+            )
+        
         self.tool_name = tool_name
         self._sm: Optional[StreamManager] = stream_manager
         self.default_timeout = default_timeout
@@ -48,7 +59,68 @@ class MCPTool:
         self._server_names = server_names or {}
         self._namespace = namespace
 
-        self._sm_lock = asyncio.Lock()
+        # Create lock only when needed (not during deserialization)
+        self._sm_lock: Optional[asyncio.Lock] = None
+
+    def _ensure_lock(self) -> asyncio.Lock:
+        """Ensure the lock exists, creating it if necessary."""
+        if self._sm_lock is None:
+            self._sm_lock = asyncio.Lock()
+        return self._sm_lock
+
+    # ------------------------------------------------------------------ #
+    # Serialization support for subprocess execution
+    # ------------------------------------------------------------------ #
+    def __getstate__(self) -> Dict[str, Any]:
+        """
+        Custom serialization for pickle support.
+        
+        Excludes non-serializable async components and stream manager.
+        The subprocess will recreate these as needed.
+        """
+        state = self.__dict__.copy()
+        
+        # Remove non-serializable items
+        state['_sm'] = None  # StreamManager will be recreated in subprocess
+        state['_sm_lock'] = None  # Lock will be recreated when needed
+        
+        # Ensure we have the necessary configuration for subprocess
+        # If no servers specified, default to the tool name (common pattern)
+        if not state.get('_servers'):
+            # Extract server name from tool_name (e.g., "get_current_time" -> "time")
+            # This is a heuristic - adjust based on your naming convention
+            if 'time' in self.tool_name.lower():
+                state['_servers'] = ['time']
+                state['_server_names'] = {0: 'time'}
+            else:
+                # Default fallback - use the tool name itself
+                state['_servers'] = [self.tool_name]
+                state['_server_names'] = {0: self.tool_name}
+        
+        # Ensure we have a config file path
+        if not state.get('_cfg_file'):
+            state['_cfg_file'] = 'server_config.json'
+        
+        logger.debug(f"Serializing MCPTool '{self.tool_name}' for subprocess with servers: {state['_servers']}")
+        return state
+    
+    def __setstate__(self, state: Dict[str, Any]) -> None:
+        """
+        Custom deserialization for pickle support.
+        
+        Restores the object state and ensures required fields are set.
+        """
+        self.__dict__.update(state)
+        
+        # Ensure critical fields exist
+        if not hasattr(self, 'tool_name') or not self.tool_name:
+            raise ValueError("Invalid MCPTool state: missing tool_name")
+        
+        # Initialize transient fields
+        self._sm = None
+        self._sm_lock = None
+        
+        logger.debug(f"Deserialized MCPTool '{self.tool_name}' in subprocess")
 
     # ------------------------------------------------------------------ #
     async def _ensure_stream_manager(self) -> StreamManager:
@@ -61,7 +133,8 @@ class MCPTool:
         if self._sm is not None:
             return self._sm
 
-        async with self._sm_lock:
+        # Use the lock, creating it if needed
+        async with self._ensure_lock():
             if self._sm is None:  # re-check inside lock
                 logger.info(
                     "Boot-strapping MCP stdio transport for '%s'", self.tool_name
@@ -139,9 +212,32 @@ class MCPTool:
 
         return result.get("content")
 
-
     # ------------------------------------------------------------------ #
     # Legacy method name support
     async def _aexecute(self, timeout: Optional[float] = None, **kwargs: Any) -> Any:
         """Legacy alias for execute() method."""
-        return await self.execute(timeout=timeout, **kwargs)
+        return await self.execute(timeout=timeout, **kwargs)
+
+    # ------------------------------------------------------------------ #
+    # Utility methods for debugging
+    # ------------------------------------------------------------------ #
+    def is_serializable(self) -> bool:
+        """Check if this tool can be serialized (for debugging)."""
+        try:
+            import pickle
+            pickle.dumps(self)
+            return True
+        except Exception:
+            return False
+    
+    def get_serialization_info(self) -> Dict[str, Any]:
+        """Get information about what would be serialized."""
+        state = self.__getstate__()
+        return {
+            "tool_name": state.get("tool_name"),
+            "namespace": state.get("_namespace"),
+            "servers": state.get("_servers"),
+            "cfg_file": state.get("_cfg_file"),
+            "has_stream_manager": state.get("_sm") is not None,
+            "serializable_size": len(str(state))
+        }

chuk_tool_processor/mcp/stream_manager.py

@@ -1,3 +1,4 @@
+#!/usr/bin/env python
 # chuk_tool_processor/mcp/stream_manager.py
 """
 StreamManager for CHUK Tool Processor.
@@ -36,6 +37,7 @@ class StreamManager:
         self.server_names: Dict[int, str] = {}
         self.all_tools: List[Dict[str, Any]] = []
         self._lock = asyncio.Lock()
+        self._close_tasks: List[asyncio.Task] = []  # Track cleanup tasks
 
     # ------------------------------------------------------------------ #
     #  factory helpers                                                   #
@@ -364,22 +366,76 @@ class StreamManager:
             return await transport.call_tool(tool_name, arguments)
         
     # ------------------------------------------------------------------ #
-    #  shutdown                                                          #
+    #  shutdown - PROPERLY FIXED VERSION                                 #
     # ------------------------------------------------------------------ #
     async def close(self) -> None:
-        tasks = [tr.close() for tr in self.transports.values()]
-        if tasks:
+        """
+        Properly close all transports with graceful handling of cancellation.
+        """
+        if not self.transports:
+            return
+        
+        # Cancel any existing close tasks
+        for task in self._close_tasks:
+            if not task.done():
+                task.cancel()
+        self._close_tasks.clear()
+        
+        # Create close tasks for all transports
+        close_tasks = []
+        for name, transport in list(self.transports.items()):
             try:
-                await asyncio.gather(*tasks)
-            except asyncio.CancelledError:  # pragma: no cover
-                pass
-            except Exception as exc:  # noqa: BLE001
-                logger.error("Error during close: %s", exc)
-
+                task = asyncio.create_task(
+                    self._close_transport(name, transport), 
+                    name=f"close_{name}"
+                )
+                close_tasks.append(task)
+                self._close_tasks.append(task)
+            except Exception as e:
+                logger.debug(f"Error creating close task for {name}: {e}")
+        
+        # Wait for all close tasks with a timeout
+        if close_tasks:
+            try:
+                # Give transports a reasonable time to close gracefully
+                await asyncio.wait_for(
+                    asyncio.gather(*close_tasks, return_exceptions=True),
+                    timeout=2.0
+                )
+            except asyncio.TimeoutError:
+                # Cancel any still-running tasks
+                for task in close_tasks:
+                    if not task.done():
+                        task.cancel()
+                # Brief wait for cancellation to take effect
+                await asyncio.gather(*close_tasks, return_exceptions=True)
+            except asyncio.CancelledError:
+                # This is expected during event loop shutdown
+                logger.debug("Close operation cancelled during shutdown")
+            except Exception as e:
+                logger.debug(f"Unexpected error during close: {e}")
+        
+        # Clean up state
+        self._cleanup_state()
+    
+    async def _close_transport(self, name: str, transport: MCPBaseTransport) -> None:
+        """Close a single transport with error handling."""
+        try:
+            await transport.close()
+            logger.debug(f"Closed transport: {name}")
+        except asyncio.CancelledError:
+            # Re-raise cancellation
+            raise
+        except Exception as e:
+            logger.debug(f"Error closing transport {name}: {e}")
+    
+    def _cleanup_state(self) -> None:
+        """Clean up internal state (synchronous)."""
         self.transports.clear()
         self.server_info.clear()
         self.tool_to_server_map.clear()
         self.all_tools.clear()
+        self._close_tasks.clear()
 
     # ------------------------------------------------------------------ #
     #  backwards-compat: streams helper                                  #

chuk_tool_processor/mcp/transport/sse_transport.py

@@ -10,6 +10,7 @@ This transport:
 5. Handles async responses via SSE message events
 
 FIXED: All hardcoded timeouts are now configurable parameters.
+FIXED: Enhanced close method to avoid cancel scope conflicts.
 """
 from __future__ import annotations
 
@@ -206,24 +207,11 @@ class SSETransport(MCPBaseTransport):
             print(f"⚠️ Failed to send notification: {e}")
 
     async def close(self) -> None:
-        """Close the transport."""
-        # Cancel any pending requests
-        for future in self._pending_requests.values():
-            if not future.done():
-                future.cancel()
-        self._pending_requests.clear()
-        
-        if self._sse_task:
-            self._sse_task.cancel()
-            with contextlib.suppress(asyncio.CancelledError):
-                await self._sse_task
-            self._sse_task = None
-
-        if self._client:
-            await self._client.aclose()
-            self._client = None
-            self.session = None
-
+        """Minimal close method with zero async operations."""
+        # Just clear references - no async operations at all
+        self._context_stack = None
+        self.read_stream = None
+        self.write_stream = None
     # ------------------------------------------------------------------ #
     # SSE Connection Handler                                             #
     # ------------------------------------------------------------------ #

chuk_tool_processor/mcp/transport/stdio_transport.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 from contextlib import AsyncExitStack
 import json
 from typing import Dict, Any, List, Optional
+import asyncio
 
 # ------------------------------------------------------------------ #
 #  Local import                                                      #
@@ -71,14 +72,11 @@ class StdioTransport(MCPBaseTransport):
             return False
 
     async def close(self) -> None:
-        if self._context_stack:
-            try:
-                await self._context_stack.__aexit__(None, None, None)
-            except Exception:
-                pass
+        """Minimal close method with zero async operations."""
+        # Just clear references - no async operations at all
+        self._context_stack = None
         self.read_stream = None
         self.write_stream = None
-        self._context_stack = None
 
     # --------------------------------------------------------------------- #
     #  Utility                                                              #
@@ -194,4 +192,4 @@ class StdioTransport(MCPBaseTransport):
             import logging
 
             logging.error(f"Error calling tool {tool_name}: {e}")
-            return {"isError": True, "error": str(e)}
+            return {"isError": True, "error": str(e)}