hanzo-mcp 0.9.0__py3-none-any.whl → 0.9.2__py3-none-any.whl

hanzo_mcp/tools/common/auto_timeout.py

@@ -0,0 +1,254 @@
+"""Universal auto-timeout and backgrounding for all MCP tools.
+
+This module provides automatic timeout and backgrounding for any MCP tool operation
+that takes longer than the configured threshold (default: 2 minutes).
+"""
+
+import asyncio
+import functools
+import json
+import time
+import uuid
+import os
+from pathlib import Path
+from typing import Any, Callable, Optional, Tuple
+from collections.abc import Awaitable
+
+from mcp.server.fastmcp import Context as MCPContext
+from .timeout_parser import parse_timeout, format_timeout
+
+
+class MCPToolTimeoutManager:
+    """Manager for MCP tool timeouts and backgrounding."""
+    
+    # Default timeout before auto-backgrounding (2 minutes)
+    DEFAULT_TIMEOUT = 120.0
+    
+    # Environment variable to configure timeout
+    TIMEOUT_ENV_VAR = "HANZO_MCP_TOOL_TIMEOUT"
+    
+    def __init__(self, process_manager: Optional[Any] = None):
+        """Initialize the timeout manager.
+
+        Args:
+            process_manager: Process manager for tracking background operations
+        """
+        if process_manager is None:
+            # Lazy import to avoid circular imports
+            try:
+                from hanzo_mcp.tools.shell.base_process import ProcessManager
+                self.process_manager = ProcessManager()
+            except ImportError:
+                # If ProcessManager is not available, disable backgrounding
+                self.process_manager = None
+        else:
+            self.process_manager = process_manager
+        
+        # Get timeout from environment or use default
+        env_timeout = os.getenv(self.TIMEOUT_ENV_VAR)
+        if env_timeout:
+            try:
+                self.timeout = parse_timeout(env_timeout)
+            except ValueError:
+                self.timeout = self.DEFAULT_TIMEOUT
+        else:
+            self.timeout = self.DEFAULT_TIMEOUT
+    
+    def _get_timeout_for_tool(self, tool_name: str) -> float:
+        """Get timeout setting for a specific tool.
+        
+        Args:
+            tool_name: Name of the tool
+            
+        Returns:
+            Timeout in seconds
+        """
+        # Check for tool-specific timeout
+        env_var = f"HANZO_MCP_{tool_name.upper()}_TIMEOUT"
+        tool_timeout = os.getenv(env_var)
+        if tool_timeout:
+            try:
+                return parse_timeout(tool_timeout)
+            except ValueError:
+                pass
+        
+        return self.timeout
+    
+    async def _background_tool_execution(
+        self,
+        tool_func: Callable,
+        tool_name: str,
+        ctx: MCPContext,
+        process_id: str,
+        log_file: Path,
+        **params: Any
+    ) -> None:
+        """Execute tool in background and log results.
+        
+        Args:
+            tool_func: The tool function to execute
+            tool_name: Name of the tool
+            ctx: MCP context
+            process_id: Process identifier
+            log_file: Log file path
+            **params: Tool parameters
+        """
+        try:
+            # Log start
+            with open(log_file, "a") as f:
+                f.write(f"=== Background execution started for {tool_name} ===\\n")
+                f.write(f"Parameters: {json.dumps(params, indent=2, default=str)}\\n")
+                f.write(f"Started at: {time.strftime('%Y-%m-%d %H:%M:%S')}\\n\\n")
+            
+            # Execute the tool
+            result = await tool_func(ctx, **params)
+            
+            # Log completion
+            with open(log_file, "a") as f:
+                f.write(f"\\n\\n=== Tool execution completed ===\\n")
+                f.write(f"Completed at: {time.strftime('%Y-%m-%d %H:%M:%S')}\\n")
+                f.write(f"Result length: {len(str(result))} characters\\n")
+                f.write("\\n=== RESULT ===\\n")
+                f.write(str(result))
+                f.write("\\n=== END RESULT ===\\n")
+            
+            # Mark as completed
+            self.process_manager.mark_completed(process_id, 0)
+            
+        except Exception as e:
+            # Log error
+            with open(log_file, "a") as f:
+                f.write(f"\\n\\n=== Tool execution failed ===\\n")
+                f.write(f"Failed at: {time.strftime('%Y-%m-%d %H:%M:%S')}\\n")
+                f.write(f"Error: {str(e)}\\n")
+                f.write(f"Error type: {type(e).__name__}\\n")
+            
+            self.process_manager.mark_completed(process_id, 1)
+
+
+def with_auto_timeout(tool_name: str, timeout_manager: Optional[MCPToolTimeoutManager] = None):
+    """Decorator to add automatic timeout and backgrounding to MCP tools.
+
+    Args:
+        tool_name: Name of the tool (for logging and process tracking)
+        timeout_manager: Optional timeout manager instance
+
+    Returns:
+        Decorator function
+    """
+    if timeout_manager is None:
+        timeout_manager = MCPToolTimeoutManager()
+
+    def decorator(func: Callable[..., Awaitable[Any]]) -> Callable[..., Awaitable[Any]]:
+        @functools.wraps(func)
+        async def wrapper(*args: Any, **params: Any) -> Any:
+            # Handle both method calls (with self) and function calls
+            # For methods: args = (self, ctx), For functions: args = (ctx,)
+            if len(args) >= 2:
+                # Method call: self, ctx, **params
+                self_or_ctx = args[0]
+                ctx = args[1]
+                call_func = lambda: func(self_or_ctx, ctx, **params)
+            elif len(args) == 1:
+                # Function call: ctx, **params
+                ctx = args[0]
+                call_func = lambda: func(ctx, **params)
+            else:
+                raise TypeError(f"Expected at least 1 argument (ctx), got {len(args)}")
+
+            # Fast path for tests - skip timeout logic
+            if os.getenv("HANZO_MCP_FAST_TESTS") == "1":
+                return await call_func()
+
+            # Get tool-specific timeout
+            tool_timeout = timeout_manager._get_timeout_for_tool(tool_name)
+
+            # Create task for the tool execution
+            tool_task = asyncio.create_task(call_func())
+
+            try:
+                # Wait for completion with timeout
+                result = await asyncio.wait_for(tool_task, timeout=tool_timeout)
+                return result
+
+            except asyncio.TimeoutError:
+                # Tool timed out - background it
+                process_id = f"{tool_name}_{uuid.uuid4().hex[:8]}"
+                log_file = timeout_manager.process_manager.create_log_file(process_id)
+
+                # Start background execution (need to reconstruct the call)
+                async def background_call():
+                    if len(args) >= 2:
+                        return await func(args[0], ctx, **params)
+                    else:
+                        return await func(ctx, **params)
+
+                asyncio.create_task(
+                    timeout_manager._background_tool_execution(
+                        background_call, tool_name, ctx, process_id, log_file, **params
+                    )
+                )
+
+                # Return backgrounding message
+                timeout_formatted = format_timeout(tool_timeout)
+                return (
+                    f"Operation automatically backgrounded after {timeout_formatted}\\n"
+                    f"Process ID: {process_id}\\n"
+                    f"Log file: {log_file}\\n\\n"
+                    f"Use 'process --action logs --id {process_id}' to view results\\n"
+                    f"Use 'process --action kill --id {process_id}' to cancel\\n\\n"
+                    f"The {tool_name} operation is continuing in the background..."
+                )
+
+        return wrapper
+    return decorator
+
+
+# Global timeout manager instance
+_global_timeout_manager = None
+
+
+def get_global_timeout_manager() -> MCPToolTimeoutManager:
+    """Get the global timeout manager instance.
+    
+    Returns:
+        Global timeout manager
+    """
+    global _global_timeout_manager
+    if _global_timeout_manager is None:
+        _global_timeout_manager = MCPToolTimeoutManager()
+    return _global_timeout_manager
+
+
+def set_global_timeout(timeout_seconds: float) -> None:
+    """Set the global timeout for all MCP tools.
+    
+    Args:
+        timeout_seconds: Timeout in seconds
+    """
+    manager = get_global_timeout_manager()
+    manager.timeout = timeout_seconds
+
+
+def set_tool_timeout(tool_name: str, timeout_seconds: float) -> None:
+    """Set timeout for a specific tool via environment variable.
+    
+    Args:
+        tool_name: Name of the tool
+        timeout_seconds: Timeout in seconds
+    """
+    env_var = f"HANZO_MCP_{tool_name.upper()}_TIMEOUT"
+    os.environ[env_var] = str(timeout_seconds)
+
+
+# Convenience decorator using global manager
+def auto_timeout(tool_name: str):
+    """Convenience decorator using the global timeout manager.
+    
+    Args:
+        tool_name: Name of the tool
+        
+    Returns:
+        Decorator function
+    """
+    return with_auto_timeout(tool_name, get_global_timeout_manager())

hanzo_mcp/tools/common/base.py

@@ -11,6 +11,8 @@ from typing import Any, Callable, final
 from collections.abc import Awaitable
 
 from mcp.server import FastMCP
+
+from hanzo_mcp.tools.common.auto_timeout import auto_timeout
 from mcp.server.fastmcp import Context as MCPContext
 
 from hanzo_mcp.tools.common.validation import (
@@ -88,6 +90,8 @@ class BaseTool(ABC):
         pass
 
     @abstractmethod
+    @auto_timeout("base")
+
     async def call(self, ctx: MCPContext, **params: Any) -> Any:
         """Execute the tool with the given parameters.
 

hanzo_mcp/tools/common/batch_tool.py

@@ -9,6 +9,8 @@ from typing import Any, Unpack, Annotated, TypedDict, final, override
 
 from pydantic import Field
 from mcp.server import FastMCP
+
+from hanzo_mcp.tools.common.auto_timeout import auto_timeout
 from mcp.server.fastmcp import Context as MCPContext
 
 from hanzo_mcp.tools.common.base import BaseTool
@@ -154,6 +156,9 @@ Not available: think,write,edit,multi_edit,notebook_edit
         self.tools = tools
 
     @override
+    @auto_timeout("batch")
+
+
     async def call(
         self,
         ctx: MCPContext,

hanzo_mcp/tools/common/config_tool.py

@@ -5,6 +5,8 @@ from typing import Any, Dict, Unpack, Optional, TypedDict, final
 
 from mcp.server.fastmcp import Context as MCPContext
 
+from hanzo_mcp.tools.common.auto_timeout import auto_timeout
+
 from hanzo_mcp.config.settings import (
     ProjectConfig,
     MCPServerConfig,
@@ -62,6 +64,9 @@ Perfect for AI-driven configuration where users can say things like:
 
 Automatically detects projects based on LLM.md files and manages .hanzo/ directories."""
 
+    @auto_timeout("config")
+
+
     async def call(
         self,
         ctx: MCPContext,

hanzo_mcp/tools/common/critic_tool.py

@@ -7,6 +7,8 @@ from typing import Unpack, Annotated, TypedDict, final, override
 
 from pydantic import Field
 from mcp.server import FastMCP
+
+from hanzo_mcp.tools.common.auto_timeout import auto_timeout
 from mcp.server.fastmcp import Context as MCPContext
 
 from hanzo_mcp.tools.common.base import BaseTool
@@ -131,6 +133,9 @@ Recommendations:
         pass
 
     @override
+    @auto_timeout("critic")
+
+
     async def call(
         self,
         ctx: MCPContext,

hanzo_mcp/tools/common/paginated_base.py

@@ -9,6 +9,8 @@ from typing import Any, Dict, Union
 
 from mcp.server.fastmcp import Context as MCPContext
 
+from hanzo_mcp.tools.common.auto_timeout import auto_timeout
+
 from hanzo_mcp.tools.common.base import BaseTool, handle_connection_errors
 from hanzo_mcp.tools.common.pagination import CursorManager
 from hanzo_mcp.tools.common.paginated_response import paginate_if_needed
@@ -44,6 +46,8 @@ class PaginatedBaseTool(BaseTool):
         pass
 
     @handle_connection_errors
+    @auto_timeout("paginated_base")
+
     async def call(self, ctx: MCPContext, **params: Any) -> Union[str, Dict[str, Any]]:
         """Execute the tool with automatic pagination support.
 

hanzo_mcp/tools/common/permissions.py

@@ -8,9 +8,6 @@ from typing import Any, TypeVar, final
 from pathlib import Path
 from collections.abc import Callable, Awaitable
 
-# Import centralized path resolution
-from hanzo_mcp.tools.common.path_utils import resolve_path
-
 # Define type variables for better type annotations
 T = TypeVar("T")
 P = TypeVar("P")
@@ -84,8 +81,7 @@ class PermissionManager:
         Args:
             path: The path to allow
         """
-        # Use centralized path resolution
-        resolved_path = Path(resolve_path(path))
+        resolved_path: Path = Path(path).resolve()
         self.allowed_paths.add(resolved_path)
 
     def remove_allowed_path(self, path: str) -> None:
@@ -94,8 +90,7 @@ class PermissionManager:
         Args:
             path: The path to remove
         """
-        # Use centralized path resolution
-        resolved_path = Path(resolve_path(path))
+        resolved_path: Path = Path(path).resolve()
         if resolved_path in self.allowed_paths:
             self.allowed_paths.remove(resolved_path)
 
@@ -105,8 +100,7 @@ class PermissionManager:
         Args:
             path: The path to exclude
         """
-        # Use centralized path resolution
-        resolved_path = Path(resolve_path(path))
+        resolved_path: Path = Path(path).resolve()
         self.excluded_paths.add(resolved_path)
 
     def add_exclusion_pattern(self, pattern: str) -> None:
@@ -118,7 +112,7 @@ class PermissionManager:
         self.excluded_patterns.append(pattern)
 
     def is_path_allowed(self, path: str) -> bool:
-        """Check if a path is allowed.
+        """Check if a path is allowed with security validation.
 
         Args:
             path: The path to check
@@ -126,8 +120,24 @@ class PermissionManager:
         Returns:
             True if the path is allowed, False otherwise
         """
-        # Use centralized path resolution
-        resolved_path = Path(resolve_path(path))
+        # Security check: Reject paths with traversal attempts
+        if ".." in str(path) or "~" in str(path):
+            return False
+        
+        try:
+            # Resolve the path (follows symlinks and makes absolute)
+            resolved_path: Path = Path(path).resolve(strict=False)
+            
+            # Security check: Ensure resolved path doesn't escape allowed directories
+            # by checking if it's actually under an allowed path after resolution
+            original_path = Path(path)
+            if original_path.is_absolute() and str(resolved_path) != str(original_path.resolve(strict=False)):
+                # Path resolution changed the path significantly, might be symlink attack
+                # Additional check: is the resolved path still under allowed paths?
+                pass  # Continue to normal checks
+        except (OSError, RuntimeError) as e:
+            # Path resolution failed, deny access
+            return False
 
         # Check exclusions first
         if self._is_path_excluded(resolved_path):
@@ -136,12 +146,28 @@ class PermissionManager:
         # Check if the path is within any allowed path
         for allowed_path in self.allowed_paths:
             try:
+                # This will raise ValueError if resolved_path is not under allowed_path
                 resolved_path.relative_to(allowed_path)
+                # Additional check: ensure no symlinks are escaping the allowed directory
+                if resolved_path.exists() and resolved_path.is_symlink():
+                    link_target = Path(os.readlink(resolved_path))
+                    if link_target.is_absolute():
+                        # Absolute symlink - check if it points within allowed paths
+                        if not any(self._is_subpath(link_target, ap) for ap in self.allowed_paths):
+                            return False
                 return True
             except ValueError:
                 continue
 
         return False
+    
+    def _is_subpath(self, child: Path, parent: Path) -> bool:
+        """Check if child is a subpath of parent."""
+        try:
+            child.resolve().relative_to(parent.resolve())
+            return True
+        except ValueError:
+            return False
 
     def _is_path_excluded(self, path: Path) -> bool:
         """Check if a path is excluded.