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

chuk_tool_processor/mcp/transport/stdio_transport.py

@@ -1,9 +1,9 @@
 # chuk_tool_processor/mcp/transport/stdio_transport.py
 from __future__ import annotations
 
-from contextlib import AsyncExitStack
 import json
 from typing import Dict, Any, List, Optional
+import logging
 
 # ------------------------------------------------------------------ #
 #  Local import                                                      #
@@ -11,187 +11,382 @@ from typing import Dict, Any, List, Optional
 from .base_transport import MCPBaseTransport
 
 # ------------------------------------------------------------------ #
-#  chuk-protocol imports                                             #
+#  New chuk-mcp API imports only                                    #
 # ------------------------------------------------------------------ #
-from chuk_mcp.mcp_client.transport.stdio.stdio_client import stdio_client
-from chuk_mcp.mcp_client.messages.initialize.send_messages import send_initialize
-from chuk_mcp.mcp_client.messages.ping.send_messages import send_ping
+from chuk_mcp.transports.stdio import stdio_client
+from chuk_mcp.transports.stdio.parameters import StdioParameters
 
-# tools
-from chuk_mcp.mcp_client.messages.tools.send_messages import (
-    send_tools_call,
+from chuk_mcp.protocol.messages import (
+    send_initialize,
+    send_ping,
     send_tools_list,
+    send_tools_call,
 )
 
-# NEW: resources & prompts
-from chuk_mcp.mcp_client.messages.resources.send_messages import (
-    send_resources_list,
-)
-from chuk_mcp.mcp_client.messages.prompts.send_messages import (
-    send_prompts_list,
-)
+# Try to import resources and prompts if available
+try:
+    from chuk_mcp.protocol.messages import (
+        send_resources_list,
+        send_resources_read,
+    )
+    HAS_RESOURCES = True
+except ImportError:
+    HAS_RESOURCES = False
+
+try:
+    from chuk_mcp.protocol.messages import (
+        send_prompts_list,
+        send_prompts_get,
+    )
+    HAS_PROMPTS = True
+except ImportError:
+    HAS_PROMPTS = False
+
+logger = logging.getLogger(__name__)
 
 
 class StdioTransport(MCPBaseTransport):
     """
-    Stdio transport for MCP communication.
+    STDIO transport for MCP communication using new chuk-mcp APIs.
     """
 
     def __init__(self, server_params):
-        self.server_params = server_params
+        """
+        Initialize STDIO transport.
+        
+        Args:
+            server_params: Either a dict with 'command' and 'args', 
+                          or a StdioParameters object
+        """
+        # Convert dict format to StdioParameters
+        if isinstance(server_params, dict):
+            self.server_params = StdioParameters(
+                command=server_params.get('command', 'python'),
+                args=server_params.get('args', [])
+            )
+        else:
+            self.server_params = server_params
+        
         self.read_stream = None
         self.write_stream = None
-        self._context_stack: Optional[AsyncExitStack] = None
+        self._client_context = None
 
     # --------------------------------------------------------------------- #
     #  Connection management                                                #
     # --------------------------------------------------------------------- #
     async def initialize(self) -> bool:
+        """Initialize the STDIO transport."""
         try:
-            self._context_stack = AsyncExitStack()
-            await self._context_stack.__aenter__()
-
-            ctx = stdio_client(self.server_params)
-            self.read_stream, self.write_stream = await self._context_stack.enter_async_context(
-                ctx
-            )
-
+            logger.info("Initializing STDIO transport...")
+            
+            # Use the new stdio_client context manager
+            self._client_context = stdio_client(self.server_params)
+            self.read_stream, self.write_stream = await self._client_context.__aenter__()
+            
+            # Send initialize message
+            logger.debug("Sending initialize message...")
             init_result = await send_initialize(self.read_stream, self.write_stream)
-            return bool(init_result)
+            
+            if init_result:
+                logger.info("STDIO transport initialized successfully")
+                return True
+            else:
+                logger.error("Initialize message failed")
+                await self._cleanup()
+                return False
 
-        except Exception as e:  # pragma: no cover
-            import logging
+        except Exception as e:
+            logger.error(f"Error initializing STDIO transport: {e}", exc_info=True)
+            await self._cleanup()
+            return False
 
-            logging.error(f"Error initializing stdio transport: {e}")
-            if self._context_stack:
+    async def close(self) -> None:
+        """Close the transport with proper error handling."""
+        try:
+            # Handle both old _context_stack and new _client_context for test compatibility
+            if hasattr(self, '_context_stack') and self._context_stack:
                 try:
                     await self._context_stack.__aexit__(None, None, None)
-                except Exception:
-                    pass
-            return False
+                    logger.debug("Context stack closed")
+                except asyncio.CancelledError:
+                    # Expected during shutdown - don't log as error
+                    logger.debug("Context stack close cancelled during shutdown")
+                except Exception as e:
+                    logger.error(f"Error closing context stack: {e}")
+            elif self._client_context:
+                try:
+                    await self._client_context.__aexit__(None, None, None)
+                    logger.debug("STDIO client context closed")
+                except asyncio.CancelledError:
+                    # Expected during shutdown - don't log as error
+                    logger.debug("Client context close cancelled during shutdown")
+                except Exception as e:
+                    logger.error(f"Error closing client context: {e}")
+        except Exception as e:
+            logger.error(f"Error during transport cleanup: {e}")
+        finally:
+            await self._cleanup()
 
-    async def close(self) -> None:
-        if self._context_stack:
-            try:
-                await self._context_stack.__aexit__(None, None, None)
-            except Exception:
-                pass
+    async def _cleanup(self) -> None:
+        """Internal cleanup method."""
+        # Clean up both old and new context attributes for test compatibility
+        if hasattr(self, '_context_stack'):
+            self._context_stack = None
+        self._client_context = None
         self.read_stream = None
         self.write_stream = None
-        self._context_stack = None
 
     # --------------------------------------------------------------------- #
-    #  Utility                                                              #
+    #  Core MCP Operations                                                  #
     # --------------------------------------------------------------------- #
     async def send_ping(self) -> bool:
+        """Send a ping."""
         if not self.read_stream or not self.write_stream:
+            logger.error("Cannot send ping: streams not available")
+            return False
+        
+        try:
+            result = await send_ping(self.read_stream, self.write_stream)
+            logger.debug(f"Ping result: {result}")
+            return bool(result)
+        except Exception as e:
+            logger.error(f"Ping failed: {e}")
             return False
-        return await send_ping(self.read_stream, self.write_stream)
 
     async def get_tools(self) -> List[Dict[str, Any]]:
+        """Get list of available tools."""
         if not self.read_stream or not self.write_stream:
+            logger.error("Cannot get tools: streams not available")
+            return []
+        
+        try:
+            tools_response = await send_tools_list(self.read_stream, self.write_stream)
+            
+            # Handle both dict response and direct tools list
+            if isinstance(tools_response, dict):
+                tools = tools_response.get("tools", [])
+            elif isinstance(tools_response, list):
+                tools = tools_response
+            else:
+                logger.warning(f"Unexpected tools response type: {type(tools_response)}")
+                tools = []
+            
+            logger.debug(f"Retrieved {len(tools)} tools")
+            return tools
+            
+        except Exception as e:
+            logger.error(f"Error getting tools: {e}")
             return []
-        tools_response = await send_tools_list(self.read_stream, self.write_stream)
-        return tools_response.get("tools", [])
 
-    # NEW ------------------------------------------------------------------ #
-    #  Resources / Prompts                                                   #
-    # --------------------------------------------------------------------- #
-    async def list_resources(self) -> Dict[str, Any]:
+    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
         """
-        Return the result of *resources/list*.  If the connection is not yet
-        initialised an empty dict is returned.
+        Execute a tool.
+        
+        Returns normalized response in format:
+        {
+            "isError": bool,
+            "content": Any,  # Result data if successful
+            "error": str     # Error message if failed
+        }
         """
         if not self.read_stream or not self.write_stream:
+            return {
+                "isError": True,
+                "error": "Transport not initialized"
+            }
+
+        try:
+            logger.debug(f"Calling tool {tool_name} with args: {arguments}")
+            
+            raw_response = await send_tools_call(
+                self.read_stream, 
+                self.write_stream, 
+                tool_name, 
+                arguments
+            )
+            
+            logger.debug(f"Tool {tool_name} raw response: {raw_response}")
+            return self._normalize_tool_response(raw_response)
+
+        except Exception as e:
+            logger.error(f"Error calling tool {tool_name}: {e}")
+            return {
+                "isError": True,
+                "error": f"Tool execution failed: {str(e)}"
+            }
+
+    def _normalize_tool_response(self, raw_response: Dict[str, Any]) -> Dict[str, Any]:
+        """Normalize tool response to consistent format."""
+        # Handle explicit error in response
+        if "error" in raw_response:
+            error_info = raw_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 (MCP standard)
+        if "result" in raw_response:
+            result = raw_response["result"]
+            
+            # If result has content, extract it
+            if isinstance(result, dict) and "content" in result:
+                return {
+                    "isError": False,
+                    "content": self._extract_content(result["content"])
+                }
+            else:
+                return {
+                    "isError": False,
+                    "content": result
+                }
+
+        # Handle direct content-based response
+        if "content" in raw_response:
+            return {
+                "isError": False,
+                "content": self._extract_content(raw_response["content"])
+            }
+
+        # Fallback: return whatever the server sent
+        return {
+            "isError": False,
+            "content": raw_response
+        }
+
+    def _extract_content(self, content_list: Any) -> Any:
+        """Extract content from MCP content list format."""
+        if not isinstance(content_list, list) or not content_list:
+            return content_list
+        
+        # Handle single content item (most common)
+        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:
+                        return json.loads(text_content)
+                    except json.JSONDecodeError:
+                        return text_content
+                else:
+                    # Non-text content (image, audio, etc.)
+                    return content_item
+        
+        # Multiple content items - return the list
+        return content_list
+
+    # --------------------------------------------------------------------- #
+    #  Resources Operations (if available)                                 #
+    # --------------------------------------------------------------------- #
+    async def list_resources(self) -> Dict[str, Any]:
+        """Get list of available resources."""
+        if not HAS_RESOURCES:
+            logger.warning("Resources not supported in current chuk-mcp version")
+            return {}
+            
+        if not self.read_stream or not self.write_stream:
+            logger.error("Cannot list resources: streams not available")
             return {}
+        
         try:
-            return await send_resources_list(self.read_stream, self.write_stream)
-        except Exception as exc:  # pragma: no cover
-            import logging
+            response = await send_resources_list(self.read_stream, self.write_stream)
+            return response if isinstance(response, dict) else {}
+        except Exception as e:
+            logger.error(f"Error listing resources: {e}")
+            return {}
 
-            logging.error(f"Error listing resources: {exc}")
+    async def read_resource(self, uri: str) -> Dict[str, Any]:
+        """Read a specific resource by URI."""
+        if not HAS_RESOURCES:
+            logger.warning("Resources not supported in current chuk-mcp version")
+            return {}
+            
+        if not self.read_stream or not self.write_stream:
+            logger.error("Cannot read resource: streams not available")
+            return {}
+        
+        try:
+            response = await send_resources_read(self.read_stream, self.write_stream, uri)
+            return response if isinstance(response, dict) else {}
+        except Exception as e:
+            logger.error(f"Error reading resource {uri}: {e}")
             return {}
 
+    # --------------------------------------------------------------------- #
+    #  Prompts Operations (if available)                                   #
+    # --------------------------------------------------------------------- #
     async def list_prompts(self) -> Dict[str, Any]:
-        """
-        Return the result of *prompts/list*.  If the connection is not yet
-        initialised an empty dict is returned.
-        """
+        """Get list of available prompts."""
+        if not HAS_PROMPTS:
+            logger.warning("Prompts not supported in current chuk-mcp version")
+            return {}
+            
         if not self.read_stream or not self.write_stream:
+            logger.error("Cannot list prompts: streams not available")
             return {}
+        
         try:
-            return await send_prompts_list(self.read_stream, self.write_stream)
-        except Exception as exc:  # pragma: no cover
-            import logging
+            response = await send_prompts_list(self.read_stream, self.write_stream)
+            return response if isinstance(response, dict) else {}
+        except Exception as e:
+            logger.error(f"Error listing prompts: {e}")
+            return {}
 
-            logging.error(f"Error listing prompts: {exc}")
+    async def get_prompt(self, name: str, arguments: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+        """Get a specific prompt by name."""
+        if not HAS_PROMPTS:
+            logger.warning("Prompts not supported in current chuk-mcp version")
+            return {}
+            
+        if not self.read_stream or not self.write_stream:
+            logger.error("Cannot get prompt: streams not available")
+            return {}
+        
+        try:
+            response = await send_prompts_get(
+                self.read_stream, 
+                self.write_stream, 
+                name, 
+                arguments or {}
+            )
+            return response if isinstance(response, dict) else {}
+        except Exception as e:
+            logger.error(f"Error getting prompt {name}: {e}")
             return {}
 
-    # OPTIONAL helper ------------------------------------------------------ #
-    def get_streams(self):
-        """
-        Expose the low-level streams so legacy callers can access them
-        directly.  The base-class' default returns an empty list; here we
-        return a single-element list when the transport is active.
-        """
+    # --------------------------------------------------------------------- #
+    #  Utility Methods                                                      #
+    # --------------------------------------------------------------------- #
+    def get_streams(self) -> List[tuple]:
+        """Get the underlying streams for advanced usage."""
         if self.read_stream and self.write_stream:
             return [(self.read_stream, self.write_stream)]
         return []
 
-    # --------------------------------------------------------------------- #
-    #  Main entry-point                                                     #
-    # --------------------------------------------------------------------- #
-    async def call_tool(
-        self, tool_name: str, arguments: Dict[str, Any]
-    ) -> Dict[str, Any]:
-        """
-        Execute *tool_name* with *arguments* and normalise the server's reply.
+    def is_connected(self) -> bool:
+        """Check if transport is connected and ready."""
+        return self.read_stream is not None and self.write_stream is not None
 
-        The echo-server often returns:
-        {
-            "content": [{"type":"text","text":"{\"message\":\"…\"}"}],
-            "isError": false
-        }
-        We unwrap that so callers just receive either a dict or a plain string.
-        """
-        if not self.read_stream or not self.write_stream:
-            return {"isError": True, "error": "Transport not initialized"}
+    async def __aenter__(self):
+        """Async context manager entry."""
+        success = await self.initialize()
+        if not success:
+            raise RuntimeError("Failed to initialize STDIO transport")
+        return self
 
-        try:
-            raw = await send_tools_call(
-                self.read_stream, self.write_stream, tool_name, arguments
-            )
-
-            # Handle explicit error wrapper
-            if "error" in raw:
-                return {
-                    "isError": True,
-                    "error": raw["error"].get("message", "Unknown error"),
-                }
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+        await self.close()
 
-            # Preferred: servers that put the answer under "result"
-            if "result" in raw:
-                return {"isError": False, "content": raw["result"]}
-
-            # Common echo-server shape: top-level "content" list
-            if "content" in raw:
-                clist = raw["content"]
-                if isinstance(clist, list) and clist:
-                    first = clist[0]
-                    if isinstance(first, dict) and first.get("type") == "text":
-                        text = first.get("text", "")
-                        # Try to parse as JSON; fall back to plain string
-                        try:
-                            parsed = json.loads(text)
-                            return {"isError": False, "content": parsed}
-                        except json.JSONDecodeError:
-                            return {"isError": False, "content": text}
-
-            # Fallback: give caller whatever the server sent
-            return {"isError": False, "content": raw}
-
-        except Exception as e:  # pragma: no cover
-            import logging
-
-            logging.error(f"Error calling tool {tool_name}: {e}")
-            return {"isError": True, "error": str(e)}
+    def __repr__(self) -> str:
+        """String representation for debugging."""
+        status = "connected" if self.is_connected() else "disconnected"
+        cmd_info = f"command={getattr(self.server_params, 'command', 'unknown')}"
+        return f"StdioTransport(status={status}, {cmd_info})"