massgen 0.0.3__py3-none-any.whl → 0.1.0__py3-none-any.whl

massgen/mcp_tools/exceptions.py

@@ -0,0 +1,338 @@
+# -*- coding: utf-8 -*-
+"""
+MCP-specific exceptions with enhanced error handling and context preservation.
+"""
+
+from datetime import datetime, timezone
+from typing import Any, Dict, Optional, Union
+
+from ..logger_config import logger
+
+
+class MCPError(Exception):
+    """
+    Base exception for MCP-related errors.
+
+    Provides structured error information and context preservation
+    with enhanced debugging capabilities.
+    """
+
+    def __init__(
+        self,
+        message: str,
+        context: Optional[Dict[str, Any]] = None,
+        error_code: Optional[str] = None,
+        timestamp: Optional[datetime] = None,
+    ):
+        super().__init__(message)
+        self.context = self._sanitize_context(context or {})
+        self.error_code = error_code
+        self.timestamp = timestamp or datetime.now(timezone.utc)
+        self.original_message = message
+
+    def _sanitize_context(self, context: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Sanitize context to remove sensitive information and ensure serializability.
+        """
+        sanitized = {}
+        sensitive_keys = {"password", "token", "secret", "key", "auth", "credential"}
+
+        for key, value in context.items():
+            if any(sensitive in key.lower() for sensitive in sensitive_keys):
+                sanitized[key] = "[REDACTED]"
+            elif isinstance(value, (str, int, float, bool, type(None))):
+                sanitized[key] = value
+            else:
+                sanitized[key] = str(value)
+
+        return sanitized
+
+    def _build_context_from_kwargs(self, base_context: Optional[Dict[str, Any]] = None, **kwargs: Any) -> Dict[str, Any]:
+        """
+        Merge base context with kwargs, ignoring None values.
+
+        Copies the provided base_context (or initializes an empty dict) and updates it
+        with key/value pairs from kwargs where the value is not None. Returns the
+        resulting context dict for use in specialized error classes.
+        """
+        context: Dict[str, Any] = dict(base_context or {})
+        for key, value in kwargs.items():
+            if value is None:
+                continue
+            context[key] = value
+        return context
+
+    def __str__(self) -> str:
+        parts = [self.original_message]
+
+        if self.error_code:
+            parts.append(f"Code: {self.error_code}")
+
+        if self.context:
+            context_items = [f"{k}={v}" for k, v in self.context.items()]
+            parts.append(f"Context: {', '.join(context_items)}")
+
+        return " | ".join(parts)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "error_type": self.__class__.__name__,
+            "message": self.original_message,
+            "error_code": self.error_code,
+            "context": self.context,
+            "timestamp": self.timestamp.isoformat(),
+        }
+
+    def log_error(self) -> None:
+        """Log the error with appropriate level and context."""
+        logger.error(
+            f"{self.__class__.__name__}: {self.original_message}",
+            extra={"mcp_error": self.to_dict()},
+        )
+
+
+class MCPConnectionError(MCPError):
+    """
+    Raised when MCP server connection fails.
+
+    Includes connection details for debugging and retry logic.
+    """
+
+    def __init__(
+        self,
+        message: str,
+        server_name: Optional[str] = None,
+        transport_type: Optional[str] = None,
+        host: Optional[str] = None,
+        port: Optional[int] = None,
+        retry_count: Optional[int] = None,
+        context: Optional[Dict[str, Any]] = None,
+        error_code: Optional[str] = None,
+    ):
+        ctx = self._build_context_from_kwargs(
+            context or {},
+            server_name=server_name,
+            transport_type=transport_type,
+            host=host,
+            port=port,
+            retry_count=retry_count,
+        )
+
+        super().__init__(message, ctx, error_code)
+
+        # Store as instance attributes for easy access
+        self.server_name = server_name
+        self.transport_type = transport_type
+        self.host = host
+        self.port = port
+        self.retry_count = retry_count
+
+
+class MCPServerError(MCPError):
+    """
+    Raised when MCP server returns an error.
+
+    Includes server error codes, HTTP status codes, and additional context.
+    """
+
+    def __init__(
+        self,
+        message: str,
+        code: Optional[Union[int, str]] = None,
+        server_name: Optional[str] = None,
+        http_status: Optional[int] = None,
+        response_data: Optional[Dict[str, Any]] = None,
+        context: Optional[Dict[str, Any]] = None,
+        error_code: Optional[str] = None,
+    ):
+        ctx = self._build_context_from_kwargs(
+            context or {},
+            server_error_code=code,
+            server_name=server_name,
+            http_status=http_status,
+            response_data=response_data,
+        )
+
+        super().__init__(message, ctx, error_code)
+
+        # Store as instance attributes
+        self.code = code
+        self.server_name = server_name
+        self.http_status = http_status
+        self.response_data = response_data
+
+
+class MCPValidationError(MCPError):
+    """
+    Raised when MCP configuration or input validation fails.
+
+    Includes detailed validation information for debugging.
+    """
+
+    def __init__(
+        self,
+        message: str,
+        field: Optional[str] = None,
+        value: Optional[Any] = None,
+        expected_type: Optional[str] = None,
+        validation_rule: Optional[str] = None,
+        context: Optional[Dict[str, Any]] = None,
+        error_code: Optional[str] = None,
+    ):
+        value_str: Optional[str] = None
+        if value is not None:
+            try:
+                value_str = str(value)
+            except Exception:
+                value_str = "[UNCONVERTIBLE]"
+            if len(value_str) > 100:
+                value_str = value_str[:100]
+
+        ctx = self._build_context_from_kwargs(
+            context or {},
+            field=field,
+            value=value_str,
+            expected_type=expected_type,
+            validation_rule=validation_rule,
+        )
+
+        super().__init__(message, ctx, error_code)
+
+        # Store as instance attributes
+        self.field = field
+        self.value = value
+        self.expected_type = expected_type
+        self.validation_rule = validation_rule
+
+
+class MCPTimeoutError(MCPError):
+    """
+    Raised when MCP operations timeout.
+
+    Includes timeout details and operation context for retry logic.
+    """
+
+    def __init__(
+        self,
+        message: str,
+        timeout_seconds: Optional[float] = None,
+        operation: Optional[str] = None,
+        elapsed_seconds: Optional[float] = None,
+        server_name: Optional[str] = None,
+        context: Optional[Dict[str, Any]] = None,
+        error_code: Optional[str] = None,
+    ):
+        ctx = self._build_context_from_kwargs(
+            context or {},
+            timeout_seconds=timeout_seconds,
+            operation=operation,
+            elapsed_seconds=elapsed_seconds,
+            server_name=server_name,
+        )
+
+        super().__init__(message, ctx, error_code)
+
+        # Store as instance attributes
+        self.timeout_seconds = timeout_seconds
+        self.operation = operation
+        self.elapsed_seconds = elapsed_seconds
+        self.server_name = server_name
+
+
+class MCPAuthenticationError(MCPError):
+    """
+    Raised when MCP authentication or authorization fails.
+
+    Includes authentication context without exposing sensitive information.
+    """
+
+    def __init__(
+        self,
+        message: str,
+        auth_type: Optional[str] = None,
+        username: Optional[str] = None,
+        server_name: Optional[str] = None,
+        permission_required: Optional[str] = None,
+        context: Optional[Dict[str, Any]] = None,
+        error_code: Optional[str] = None,
+    ):
+        ctx = self._build_context_from_kwargs(
+            context or {},
+            auth_type=auth_type,
+            username=username,
+            server_name=server_name,
+            permission_required=permission_required,
+        )
+
+        super().__init__(message, ctx, error_code)
+
+        # Store as instance attributes
+        self.auth_type = auth_type
+        self.username = username
+        self.server_name = server_name
+        self.permission_required = permission_required
+
+
+class MCPConfigurationError(MCPError):
+    """
+    Raised when MCP configuration is invalid or missing.
+
+    Includes configuration details for troubleshooting.
+    """
+
+    def __init__(
+        self,
+        message: str,
+        config_file: Optional[str] = None,
+        config_section: Optional[str] = None,
+        missing_keys: Optional[list] = None,
+        context: Optional[Dict[str, Any]] = None,
+        error_code: Optional[str] = None,
+    ):
+        ctx = self._build_context_from_kwargs(
+            context or {},
+            config_file=config_file,
+            config_section=config_section,
+            missing_keys=missing_keys,
+        )
+
+        super().__init__(message, ctx, error_code)
+
+        # Store as instance attributes
+        self.config_file = config_file
+        self.config_section = config_section
+        self.missing_keys = missing_keys
+
+
+class MCPResourceError(MCPError):
+    """
+    Raised when MCP resource operations fail.
+
+    Includes resource details and operation context.
+    """
+
+    def __init__(
+        self,
+        message: str,
+        resource_type: Optional[str] = None,
+        resource_id: Optional[str] = None,
+        operation: Optional[str] = None,
+        server_name: Optional[str] = None,
+        context: Optional[Dict[str, Any]] = None,
+        error_code: Optional[str] = None,
+    ):
+        ctx = self._build_context_from_kwargs(
+            context or {},
+            resource_type=resource_type,
+            resource_id=resource_id,
+            operation=operation,
+            server_name=server_name,
+        )
+
+        super().__init__(message, ctx, error_code)
+
+        # Store as instance attributes
+        self.resource_type = resource_type
+        self.resource_id = resource_id
+        self.operation = operation
+        self.server_name = server_name

massgen/mcp_tools/hooks.py

@@ -0,0 +1,212 @@
+# -*- coding: utf-8 -*-
+"""
+Hook system for MCP tool call interception.
+
+This module provides the infrastructure for intercepting MCP tool calls
+across different backend architectures:
+
+1. Function-based backends (OpenAI, Claude, etc.) - use FunctionHook
+2. Session-based backends (Gemini) - use PermissionClientSession
+
+The actual permission logic is implemented in filesystem_manager.py
+"""
+
+from abc import ABC, abstractmethod
+from datetime import timedelta
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+from ..logger_config import logger
+
+# MCP imports for session-based backends
+try:
+    from mcp import ClientSession, types
+    from mcp.client.session import ProgressFnT
+
+    MCP_AVAILABLE = True
+except ImportError:
+    MCP_AVAILABLE = False
+    ClientSession = object
+    types = None
+    ProgressFnT = None
+
+
+class HookType(Enum):
+    """Types of function call hooks."""
+
+    PRE_CALL = "pre_call"
+    # Future: POST_CALL = "post_call"
+
+
+class HookResult:
+    """Result of a hook execution."""
+
+    def __init__(self, allowed: bool, metadata: Optional[Dict[str, Any]] = None, modified_args: Optional[str] = None):
+        self.allowed = allowed
+        self.metadata = metadata or {}
+        self.modified_args = modified_args
+
+
+class FunctionHook(ABC):
+    """Base class for function call hooks."""
+
+    def __init__(self, name: str):
+        self.name = name
+
+    @abstractmethod
+    async def execute(self, function_name: str, arguments: str, context: Optional[Dict[str, Any]] = None, **kwargs) -> HookResult:
+        """
+        Execute the hook.
+
+        Args:
+            function_name: Name of the function being called
+            arguments: JSON string of arguments
+            context: Additional context (backend, timestamp, etc.)
+
+        Returns:
+            HookResult with allowed flag and optional modifications
+        """
+
+
+class FunctionHookManager:
+    """Manages registration and execution of function hooks."""
+
+    def __init__(self):
+        self._hooks: Dict[HookType, List[FunctionHook]] = {hook_type: [] for hook_type in HookType}
+        self._global_hooks: Dict[HookType, List[FunctionHook]] = {hook_type: [] for hook_type in HookType}
+
+    def register_hook(self, function_name: str, hook_type: HookType, hook: FunctionHook):
+        """Register a hook for a specific function."""
+        if function_name not in self._hooks:
+            self._hooks[function_name] = {hook_type: [] for hook_type in HookType}
+
+        if hook_type not in self._hooks[function_name]:
+            self._hooks[function_name][hook_type] = []
+
+        self._hooks[function_name][hook_type].append(hook)
+
+    def register_global_hook(self, hook_type: HookType, hook: FunctionHook):
+        """Register a hook that applies to all functions."""
+        self._global_hooks[hook_type].append(hook)
+
+    def get_hooks_for_function(self, function_name: str) -> Dict[HookType, List[FunctionHook]]:
+        """Get all hooks (function-specific + global) for a function."""
+        result = {hook_type: [] for hook_type in HookType}
+
+        # Add global hooks first
+        for hook_type in HookType:
+            result[hook_type].extend(self._global_hooks[hook_type])
+
+        # Add function-specific hooks
+        if function_name in self._hooks:
+            for hook_type in HookType:
+                if hook_type in self._hooks[function_name]:
+                    result[hook_type].extend(self._hooks[function_name][hook_type])
+
+        return result
+
+    def clear_hooks(self):
+        """Clear all registered hooks."""
+        self._hooks.clear()
+        self._global_hooks = {hook_type: [] for hook_type in HookType}
+
+
+class PermissionClientSession(ClientSession):
+    """
+    ClientSession subclass that intercepts tool calls to apply permission hooks.
+
+    This inherits from ClientSession instead of wrapping it, which ensures
+    compatibility with SDK type checking and attribute access.
+    """
+
+    def __init__(self, wrapped_session: ClientSession, permission_manager):
+        """
+        Initialize by copying state from an existing ClientSession.
+
+        Args:
+            wrapped_session: The actual ClientSession to copy state from
+            permission_manager: Object with pre_tool_use_hook method for validation
+        """
+        # Store the permission manager
+        self._permission_manager = permission_manager
+
+        # Copy all attributes from the wrapped session to this instance
+        # This is a bit hacky but necessary to preserve the session state
+        self.__dict__.update(wrapped_session.__dict__)
+
+        logger.debug(f"[PermissionClientSession] Created permission session from {id(wrapped_session)}")
+
+    async def call_tool(
+        self,
+        name: str,
+        arguments: dict[str, Any] | None = None,
+        read_timeout_seconds: timedelta | None = None,
+        progress_callback: ProgressFnT | None = None,
+    ) -> types.CallToolResult:
+        """
+        Override call_tool to apply permission hooks before calling the actual tool.
+        """
+        tool_args = arguments or {}
+
+        # Log tool call for debugging
+        logger.debug(f"[PermissionClientSession] Intercepted tool call: {name} with args: {tool_args}")
+
+        # Apply permission hook if available
+        if self._permission_manager and hasattr(self._permission_manager, "pre_tool_use_hook"):
+            try:
+                allowed, reason = await self._permission_manager.pre_tool_use_hook(name, tool_args)
+
+                if not allowed:
+                    error_msg = f"Permission denied for tool '{name}'"
+                    if reason:
+                        error_msg += f": {reason}"
+                    logger.warning(f"🚫 [PermissionClientSession] {error_msg}")
+
+                    # Return an error result instead of calling the tool
+                    return types.CallToolResult(content=[types.TextContent(type="text", text=f"Error: {error_msg}")], isError=True)
+                else:
+                    logger.debug(f"[PermissionClientSession] Tool '{name}' permission check passed")
+
+            except Exception as e:
+                logger.error(f"[PermissionClientSession] Error in permission hook: {e}")
+                # Continue with the call if hook fails - don't break functionality
+
+        # Call the parent's call_tool method
+        try:
+            result = await super().call_tool(name=name, arguments=arguments, read_timeout_seconds=read_timeout_seconds, progress_callback=progress_callback)
+            logger.debug(f"[PermissionClientSession] Tool '{name}' completed successfully")
+            return result
+        except Exception as e:
+            logger.error(f"[PermissionClientSession] Tool '{name}' failed: {e}")
+            raise
+
+
+def convert_sessions_to_permission_sessions(sessions: List[ClientSession], permission_manager) -> List[PermissionClientSession]:
+    """
+    Convert a list of ClientSession objects to PermissionClientSession subclasses.
+
+    Args:
+        sessions: List of ClientSession objects to convert
+        permission_manager: Object with pre_tool_use_hook method
+
+    Returns:
+        List of PermissionClientSession objects that apply permission hooks
+    """
+    logger.debug(f"[PermissionClientSession] Converting {len(sessions)} sessions to permission sessions")
+    converted = []
+    for session in sessions:
+        # Create a new PermissionClientSession that inherits from ClientSession
+        perm_session = PermissionClientSession(session, permission_manager)
+        converted.append(perm_session)
+    logger.debug(f"[PermissionClientSession] Successfully converted {len(converted)} sessions")
+    return converted
+
+
+__all__ = [
+    "HookType",
+    "HookResult",
+    "FunctionHook",
+    "FunctionHookManager",
+    "PermissionClientSession",
+    "convert_sessions_to_permission_sessions",
+]