mcp-proxy-adapter 3.1.6__py3-none-any.whl → 4.1.0__py3-none-any.whl

mcp_proxy_adapter/commands/command_registry.py

@@ -41,6 +41,7 @@ class CommandRegistry:
         """
         self._commands: Dict[str, Type[Command]] = {}
         self._instances: Dict[str, Command] = {}
+        self._custom_commands: Dict[str, Type[Command]] = {}  # Custom commands with priority
     
     def register(self, command: Union[Type[Command], Command]) -> None:
         """
@@ -196,7 +197,7 @@ class CommandRegistry:
         Raises:
             NotFoundError: If command is not found.
         """
-        command_class = self.get_command(command_name)
+        command_class = self.get_command_with_priority(command_name)
         
         return {
             "name": command_name,
@@ -219,7 +220,7 @@ class CommandRegistry:
         Raises:
             NotFoundError: If command is not found
         """
-        command_class = self.get_command(command_name)
+        command_class = self.get_command_with_priority(command_name)
         return command_class.get_metadata()
     
     def get_all_metadata(self) -> Dict[str, Dict[str, Any]]:
@@ -230,8 +231,13 @@ class CommandRegistry:
             Dict with command names as keys and metadata as values
         """
         metadata = {}
-        for name, command_class in self._commands.items():
+        # Add custom commands first (they have priority)
+        for name, command_class in self._custom_commands.items():
             metadata[name] = command_class.get_metadata()
+        # Add built-in commands (custom commands will override if same name)
+        for name, command_class in self._commands.items():
+            if name not in self._custom_commands:  # Only add if not overridden by custom
+                metadata[name] = command_class.get_metadata()
         return metadata
     
     def get_all_commands_info(self) -> Dict[str, Dict[str, Any]]:
@@ -242,19 +248,29 @@ class CommandRegistry:
             Dictionary with information about all commands.
         """
         commands_info = {}
-        for name in self._commands:
+        # Add custom commands first (they have priority)
+        for name in self._custom_commands:
             commands_info[name] = self.get_command_info(name)
+        # Add built-in commands (custom commands will override if same name)
+        for name in self._commands:
+            if name not in self._custom_commands:  # Only add if not overridden by custom
+                commands_info[name] = self.get_command_info(name)
         return commands_info
     
-    def discover_commands(self, package_path: str = "mcp_proxy_adapter.commands") -> None:
+    def discover_commands(self, package_path: str = "mcp_proxy_adapter.commands") -> int:
         """
         Automatically discovers and registers commands in the specified package.
 
         Args:
             package_path: Path to package with commands.
+            
+        Returns:
+            Number of commands discovered and registered.
         """
         logger.info(f"Discovering commands in package: {package_path}")
         
+        commands_discovered = 0
+        
         try:
             package = importlib.import_module(package_path)
             package_dir = os.path.dirname(package.__file__ or "")
@@ -262,7 +278,7 @@ class CommandRegistry:
             for _, module_name, is_pkg in pkgutil.iter_modules([package_dir]):
                 if is_pkg:
                     # Recursively traverse subpackages
-                    self.discover_commands(f"{package_path}.{module_name}")
+                    commands_discovered += self.discover_commands(f"{package_path}.{module_name}")
                 elif module_name.endswith("_command"):
                     # Import only command modules
                     module_path = f"{package_path}.{module_name}"
@@ -286,6 +302,8 @@ class CommandRegistry:
                                 # Register the command only if it doesn't exist
                                 if not self.command_exists(command_name):
                                     self.register(cast(Type[Command], obj))
+                                    commands_discovered += 1
+                                    logger.debug(f"Registered command: {command_name}")
                                 else:
                                     logger.debug(f"Command '{command_name}' is already registered, skipping")
                     except ValueError as e:
@@ -296,13 +314,241 @@ class CommandRegistry:
         except Exception as e:
             logger.error(f"Error discovering commands: {e}")
             
+        return commands_discovered
+            
+    def register_custom_command(self, command: Union[Type[Command], Command]) -> None:
+        """
+        Register a custom command with priority over built-in commands.
+        
+        Args:
+            command: Command class or instance to register.
+            
+        Raises:
+            ValueError: If command with the same name is already registered.
+        """
+        # Determine if this is a class or an instance
+        if isinstance(command, type) and issubclass(command, Command):
+            command_class = command
+            command_instance = None
+        elif isinstance(command, Command):
+            command_class = command.__class__
+            command_instance = command
+        else:
+            raise ValueError(f"Invalid command type: {type(command)}. Expected Command class or instance.")
+            
+        # Get command name
+        if not hasattr(command_class, "name") or not command_class.name:
+            # Use class name if name attribute is not set
+            command_name = command_class.__name__.lower()
+            if command_name.endswith("command"):
+                command_name = command_name[:-7]  # Remove "command" suffix
+        else:
+            command_name = command_class.name
+            
+        if command_name in self._custom_commands:
+            logger.debug(f"Custom command '{command_name}' is already registered, skipping")
+            raise ValueError(f"Custom command '{command_name}' is already registered")
+            
+        logger.debug(f"Registering custom command: {command_name}")
+        self._custom_commands[command_name] = command_class
+        
+        # Store instance if provided
+        if command_instance:
+            logger.debug(f"Storing custom instance for command: {command_name}")
+            self._instances[command_name] = command_instance
+    
+    def unregister_custom_command(self, command_name: str) -> None:
+        """
+        Remove custom command from registry.
+        
+        Args:
+            command_name: Command name to remove.
+            
+        Raises:
+            NotFoundError: If command is not found.
+        """
+        if command_name not in self._custom_commands:
+            raise NotFoundError(f"Custom command '{command_name}' not found")
+            
+        logger.debug(f"Unregistering custom command: {command_name}")
+        del self._custom_commands[command_name]
+        
+        # Also remove from instances if present
+        if command_name in self._instances:
+            del self._instances[command_name]
+    
+    def custom_command_exists(self, command_name: str) -> bool:
+        """
+        Check if custom command exists.
+        
+        Args:
+            command_name: Command name to check.
+            
+        Returns:
+            True if custom command exists, False otherwise.
+        """
+        return command_name in self._custom_commands
+    
+    def get_custom_command(self, command_name: str) -> Type[Command]:
+        """
+        Get custom command class.
+        
+        Args:
+            command_name: Command name.
+            
+        Returns:
+            Command class.
+            
+        Raises:
+            NotFoundError: If command is not found.
+        """
+        if command_name not in self._custom_commands:
+            raise NotFoundError(f"Custom command '{command_name}' not found")
+        return self._custom_commands[command_name]
+    
+    def get_all_custom_commands(self) -> Dict[str, Type[Command]]:
+        """
+        Get all custom commands.
+        
+        Returns:
+            Dictionary with custom command names as keys and classes as values.
+        """
+        return self._custom_commands.copy()
+    
+    def get_priority_command(self, command_name: str) -> Optional[Type[Command]]:
+        """
+        Get command with priority (custom commands first, then built-in).
+        
+        Args:
+            command_name: Command name.
+            
+        Returns:
+            Command class if found, None otherwise.
+        """
+        # First check custom commands
+        if command_name in self._custom_commands:
+            return self._custom_commands[command_name]
+        
+        # Then check built-in commands
+        if command_name in self._commands:
+            return self._commands[command_name]
+        
+        return None
+    
+    def command_exists_with_priority(self, command_name: str) -> bool:
+        """
+        Check if command exists (custom or built-in).
+        
+        Args:
+            command_name: Command name to check.
+            
+        Returns:
+            True if command exists, False otherwise.
+        """
+        return (command_name in self._custom_commands or 
+                command_name in self._commands)
+    
+    def get_command_with_priority(self, command_name: str) -> Type[Command]:
+        """
+        Get command with priority (custom commands first, then built-in).
+        
+        Args:
+            command_name: Command name.
+            
+        Returns:
+            Command class.
+            
+        Raises:
+            NotFoundError: If command is not found.
+        """
+        # First check custom commands
+        if command_name in self._custom_commands:
+            return self._custom_commands[command_name]
+        
+        # Then check built-in commands
+        if command_name in self._commands:
+            return self._commands[command_name]
+        
+        raise NotFoundError(f"Command '{command_name}' not found")
+    
     def clear(self) -> None:
         """
-        Clears command registry.
+        Clear all registered commands.
         """
-        logger.debug("Clearing command registry")
+        logger.debug("Clearing all registered commands")
         self._commands.clear()
         self._instances.clear()
+        self._custom_commands.clear()
+    
+    def reload_config_and_commands(self, package_path: str = "mcp_proxy_adapter.commands") -> Dict[str, Any]:
+        """
+        Reload configuration and rediscover commands.
+        
+        Args:
+            package_path: Path to package with commands.
+            
+        Returns:
+            Dictionary with reload information including:
+            - config_reloaded: Whether config was reloaded
+            - commands_discovered: Number of commands discovered
+            - custom_commands_preserved: Number of custom commands preserved
+            - total_commands: Total number of commands after reload
+        """
+        logger.info("🔄 Starting configuration and commands reload...")
+        
+        # Store current custom commands
+        custom_commands_backup = self._custom_commands.copy()
+        
+        # Reload configuration
+        try:
+            from mcp_proxy_adapter.config import config
+            config.load_config()
+            config_reloaded = True
+            logger.info("✅ Configuration reloaded successfully")
+        except Exception as e:
+            logger.error(f"❌ Failed to reload configuration: {e}")
+            config_reloaded = False
+        
+        # Reinitialize logging with new configuration
+        try:
+            from mcp_proxy_adapter.core.logging import setup_logging
+            setup_logging()
+            logger.info("✅ Logging reinitialized with new configuration")
+        except Exception as e:
+            logger.error(f"❌ Failed to reinitialize logging: {e}")
+        
+        # Clear all commands except custom ones
+        self._commands.clear()
+        self._instances.clear()
+        
+        # Restore custom commands
+        self._custom_commands = custom_commands_backup
+        custom_commands_preserved = len(custom_commands_backup)
+        
+        # Rediscover commands
+        try:
+            commands_discovered = self.discover_commands(package_path)
+            logger.info(f"✅ Rediscovered {commands_discovered} commands")
+        except Exception as e:
+            logger.error(f"❌ Failed to rediscover commands: {e}")
+            commands_discovered = 0
+        
+        # Get final counts
+        total_commands = len(self._commands)
+        built_in_commands = total_commands - custom_commands_preserved
+        custom_commands = custom_commands_preserved
+        
+        result = {
+            "config_reloaded": config_reloaded,
+            "commands_discovered": commands_discovered,
+            "custom_commands_preserved": custom_commands_preserved,
+            "total_commands": total_commands,
+            "built_in_commands": built_in_commands,
+            "custom_commands": custom_commands
+        }
+        
+        logger.info(f"🔄 Reload completed: {result}")
+        return result
 
 
 # Global command registry instance

mcp_proxy_adapter/commands/hooks.py

@@ -0,0 +1,260 @@
+"""
+Module for command execution hooks.
+
+This module provides a hook system that allows intercepting command execution
+before and after the actual command runs.
+"""
+
+from typing import Any, Callable, Dict, List, Optional, Union
+from dataclasses import dataclass
+from enum import Enum
+
+from mcp_proxy_adapter.core.logging import logger
+
+
+class HookType(Enum):
+    """Types of hooks."""
+    BEFORE_EXECUTION = "before_execution"
+    AFTER_EXECUTION = "after_execution"
+
+
+@dataclass
+class HookContext:
+    """Context passed to hook functions."""
+    command_name: str
+    params: Dict[str, Any]
+    hook_type: HookType
+    standard_processing: bool = True
+    result: Optional[Any] = None
+
+
+class CommandHooks:
+    """
+    Manages command execution hooks.
+    """
+    
+    def __init__(self):
+        """Initialize hooks manager."""
+        self._before_hooks: Dict[str, List[Callable]] = {}
+        self._after_hooks: Dict[str, List[Callable]] = {}
+        self._global_before_hooks: List[Callable] = []
+        self._global_after_hooks: List[Callable] = []
+    
+    def register_before_hook(self, command_name: str, hook: Callable[[HookContext], None]) -> None:
+        """
+        Register a hook to be executed before command execution.
+        
+        Args:
+            command_name: Name of the command to hook into
+            hook: Hook function that takes HookContext as parameter
+        """
+        if command_name not in self._before_hooks:
+            self._before_hooks[command_name] = []
+        self._before_hooks[command_name].append(hook)
+        logger.debug(f"Registered before hook for command: {command_name}")
+    
+    def register_after_hook(self, command_name: str, hook: Callable[[HookContext], None]) -> None:
+        """
+        Register a hook to be executed after command execution.
+        
+        Args:
+            command_name: Name of the command to hook into
+            hook: Hook function that takes HookContext as parameter
+        """
+        if command_name not in self._after_hooks:
+            self._after_hooks[command_name] = []
+        self._after_hooks[command_name].append(hook)
+        logger.debug(f"Registered after hook for command: {command_name}")
+    
+    def register_global_before_hook(self, hook: Callable[[HookContext], None]) -> None:
+        """
+        Register a global hook to be executed before any command.
+        
+        Args:
+            hook: Hook function that takes HookContext as parameter
+        """
+        self._global_before_hooks.append(hook)
+        logger.debug("Registered global before hook")
+    
+    def register_global_after_hook(self, hook: Callable[[HookContext], None]) -> None:
+        """
+        Register a global hook to be executed after any command.
+        
+        Args:
+            hook: Hook function that takes HookContext as parameter
+        """
+        self._global_after_hooks.append(hook)
+        logger.debug("Registered global after hook")
+    
+    def unregister_before_hook(self, command_name: str, hook: Callable[[HookContext], None]) -> None:
+        """
+        Unregister a before hook for a specific command.
+        
+        Args:
+            command_name: Name of the command
+            hook: Hook function to unregister
+        """
+        if command_name in self._before_hooks:
+            try:
+                self._before_hooks[command_name].remove(hook)
+                logger.debug(f"Unregistered before hook for command: {command_name}")
+                # Remove the command key if no hooks remain
+                if not self._before_hooks[command_name]:
+                    del self._before_hooks[command_name]
+            except ValueError:
+                logger.warning(f"Hook not found for command: {command_name}")
+    
+    def unregister_after_hook(self, command_name: str, hook: Callable[[HookContext], None]) -> None:
+        """
+        Unregister an after hook for a specific command.
+        
+        Args:
+            command_name: Name of the command
+            hook: Hook function to unregister
+        """
+        if command_name in self._after_hooks:
+            try:
+                self._after_hooks[command_name].remove(hook)
+                logger.debug(f"Unregistered after hook for command: {command_name}")
+                # Remove the command key if no hooks remain
+                if not self._after_hooks[command_name]:
+                    del self._after_hooks[command_name]
+            except ValueError:
+                logger.warning(f"Hook not found for command: {command_name}")
+    
+    def unregister_global_before_hook(self, hook: Callable[[HookContext], None]) -> None:
+        """
+        Unregister a global before hook.
+        
+        Args:
+            hook: Hook function to unregister
+        """
+        try:
+            self._global_before_hooks.remove(hook)
+            logger.debug("Unregistered global before hook")
+        except ValueError:
+            logger.warning("Global before hook not found")
+    
+    def unregister_global_after_hook(self, hook: Callable[[HookContext], None]) -> None:
+        """
+        Unregister a global after hook.
+        
+        Args:
+            hook: Hook function to unregister
+        """
+        try:
+            self._global_after_hooks.remove(hook)
+            logger.debug("Unregistered global after hook")
+        except ValueError:
+            logger.warning("Global after hook not found")
+    
+    def execute_before_hooks(self, command_name: str, params: Dict[str, Any]) -> HookContext:
+        """
+        Execute all before hooks for a command.
+        
+        Args:
+            command_name: Name of the command
+            params: Command parameters
+            
+        Returns:
+            HookContext with execution results
+        """
+        context = HookContext(
+            command_name=command_name,
+            params=params,
+            hook_type=HookType.BEFORE_EXECUTION
+        )
+        
+        # Execute global before hooks
+        for hook in self._global_before_hooks:
+            try:
+                hook(context)
+            except Exception as e:
+                logger.error(f"Error in global before hook for command {command_name}: {e}")
+        
+        # Execute command-specific before hooks
+        if command_name in self._before_hooks:
+            for hook in self._before_hooks[command_name]:
+                try:
+                    hook(context)
+                except Exception as e:
+                    logger.error(f"Error in before hook for command {command_name}: {e}")
+        
+        return context
+    
+    def execute_after_hooks(self, command_name: str, params: Dict[str, Any], result: Any) -> HookContext:
+        """
+        Execute all after hooks for a command.
+        
+        Args:
+            command_name: Name of the command
+            params: Command parameters
+            result: Command execution result
+            
+        Returns:
+            HookContext with execution results
+        """
+        context = HookContext(
+            command_name=command_name,
+            params=params,
+            hook_type=HookType.AFTER_EXECUTION,
+            result=result
+        )
+        
+        # Execute command-specific after hooks
+        if command_name in self._after_hooks:
+            for hook in self._after_hooks[command_name]:
+                try:
+                    hook(context)
+                except Exception as e:
+                    logger.error(f"Error in after hook for command {command_name}: {e}")
+        
+        # Execute global after hooks
+        for hook in self._global_after_hooks:
+            try:
+                hook(context)
+            except Exception as e:
+                logger.error(f"Error in global after hook for command {command_name}: {e}")
+        
+        return context
+    
+    def clear_hooks(self, command_name: Optional[str] = None) -> None:
+        """
+        Clear all hooks or hooks for a specific command.
+        
+        Args:
+            command_name: If provided, clear hooks only for this command.
+                          If None, clear all hooks.
+        """
+        if command_name is None:
+            # Clear all hooks
+            self._before_hooks.clear()
+            self._after_hooks.clear()
+            self._global_before_hooks.clear()
+            self._global_after_hooks.clear()
+            logger.debug("Cleared all hooks")
+        else:
+            # Clear hooks for specific command
+            if command_name in self._before_hooks:
+                del self._before_hooks[command_name]
+            if command_name in self._after_hooks:
+                del self._after_hooks[command_name]
+            logger.debug(f"Cleared hooks for command: {command_name}")
+    
+    def get_hook_info(self) -> Dict[str, Any]:
+        """
+        Get information about registered hooks.
+        
+        Returns:
+            Dictionary with hook information
+        """
+        return {
+            "before_hooks": {cmd: len(hooks) for cmd, hooks in self._before_hooks.items()},
+            "after_hooks": {cmd: len(hooks) for cmd, hooks in self._after_hooks.items()},
+            "global_before_hooks": len(self._global_before_hooks),
+            "global_after_hooks": len(self._global_after_hooks)
+        }
+
+
+# Global hooks instance
+hooks = CommandHooks()