mcp-proxy-adapter 4.1.1__py3-none-any.whl → 6.0.1__py3-none-any.whl

mcp_proxy_adapter/commands/hooks.py

@@ -1,260 +1,293 @@
 """
-Module for command execution hooks.
+Module for command registration hooks.
 
-This module provides a hook system that allows intercepting command execution
-before and after the actual command runs.
+This module provides a hook system for registering custom commands
+that will be called during system initialization.
 """
 
-from typing import Any, Callable, Dict, List, Optional, Union
-from dataclasses import dataclass
+from typing import Callable, List, Optional, Dict, Any, Union
 from enum import Enum
-
+from dataclasses import dataclass
 from mcp_proxy_adapter.core.logging import logger
 
 
 class HookType(Enum):
-    """Types of hooks."""
+    """Types of hooks that can be registered."""
+    CUSTOM_COMMANDS = "custom_commands"
+    BEFORE_INIT = "before_init"
+    AFTER_INIT = "after_init"
+    BEFORE_COMMAND = "before_command"
+    AFTER_COMMAND = "after_command"
     BEFORE_EXECUTION = "before_execution"
     AFTER_EXECUTION = "after_execution"
 
 
 @dataclass
 class HookContext:
-    """Context passed to hook functions."""
-    command_name: str
-    params: Dict[str, Any]
+    """Context object passed to hook functions."""
     hook_type: HookType
-    standard_processing: bool = True
+    command_name: Optional[str] = None
+    params: Optional[Dict[str, Any]] = None
     result: Optional[Any] = None
+    registry: Optional[Any] = None
+    metadata: Optional[Dict[str, Any]] = None
+    standard_processing: bool = True
+    
+    def __post_init__(self):
+        """Initialize metadata if not provided."""
+        if self.metadata is None:
+            self.metadata = {}
 
 
 class CommandHooks:
     """
-    Manages command execution hooks.
+    Hook system for command registration.
     """
     
     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] = []
+        """
+        Initialize command hooks.
+        """
+        self._custom_commands_hooks: List[Callable] = []
+        self._before_init_hooks: List[Callable] = []
+        self._after_init_hooks: List[Callable] = []
+        self._before_command_hooks: List[Callable] = []
+        self._after_command_hooks: List[Callable] = []
     
-    def register_before_hook(self, command_name: str, hook: Callable[[HookContext], None]) -> None:
+    def register_custom_commands_hook(self, hook_func: Callable) -> None:
         """
-        Register a hook to be executed before command execution.
+        Register a hook function for custom commands registration.
         
         Args:
-            command_name: Name of the command to hook into
-            hook: Hook function that takes HookContext as parameter
+            hook_func: Function that registers custom commands.
+                      Should accept registry 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}")
+        self._custom_commands_hooks.append(hook_func)
+        logger.debug(f"Registered custom commands hook: {hook_func.__name__}")
     
-    def register_after_hook(self, command_name: str, hook: Callable[[HookContext], None]) -> None:
+    def register_before_init_hook(self, hook_func: Callable) -> None:
         """
-        Register a hook to be executed after command execution.
+        Register a hook function to be called before system initialization.
         
         Args:
-            command_name: Name of the command to hook into
-            hook: Hook function that takes HookContext as parameter
+            hook_func: Function to call before initialization.
         """
-        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}")
+        self._before_init_hooks.append(hook_func)
+        logger.debug(f"Registered before init hook: {hook_func.__name__}")
     
-    def register_global_before_hook(self, hook: Callable[[HookContext], None]) -> None:
+    def register_after_init_hook(self, hook_func: Callable) -> None:
         """
-        Register a global hook to be executed before any command.
+        Register a hook function to be called after system initialization.
         
         Args:
-            hook: Hook function that takes HookContext as parameter
+            hook_func: Function to call after initialization.
         """
-        self._global_before_hooks.append(hook)
-        logger.debug("Registered global before hook")
+        self._after_init_hooks.append(hook_func)
+        logger.debug(f"Registered after init hook: {hook_func.__name__}")
     
-    def register_global_after_hook(self, hook: Callable[[HookContext], None]) -> None:
+    def register_before_command_hook(self, hook_func: Callable) -> None:
         """
-        Register a global hook to be executed after any command.
+        Register a hook function to be called before command execution.
         
         Args:
-            hook: Hook function that takes HookContext as parameter
+            hook_func: Function to call before command execution.
+                      Should accept command_name and params as parameters.
         """
-        self._global_after_hooks.append(hook)
-        logger.debug("Registered global after hook")
+        self._before_command_hooks.append(hook_func)
+        logger.debug(f"Registered before command hook: {hook_func.__name__}")
     
-    def unregister_before_hook(self, command_name: str, hook: Callable[[HookContext], None]) -> None:
+    def register_after_command_hook(self, hook_func: Callable) -> None:
         """
-        Unregister a before hook for a specific command.
+        Register a hook function to be called after command execution.
         
         Args:
-            command_name: Name of the command
-            hook: Hook function to unregister
+            hook_func: Function to call after command execution.
+                      Should accept command_name, params, and result as parameters.
         """
-        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}")
+        self._after_command_hooks.append(hook_func)
+        logger.debug(f"Registered after command hook: {hook_func.__name__}")
     
-    def unregister_after_hook(self, command_name: str, hook: Callable[[HookContext], None]) -> None:
+    def execute_custom_commands_hooks(self, registry) -> int:
         """
-        Unregister an after hook for a specific command.
+        Execute all registered custom commands hooks.
         
         Args:
-            command_name: Name of the command
-            hook: Hook function to unregister
+            registry: Command registry instance.
+            
+        Returns:
+            Number of hooks executed.
         """
-        if command_name in self._after_hooks:
+        logger.debug("Executing custom commands hooks...")
+        hooks_executed = 0
+        
+        for hook_func in self._custom_commands_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}")
+                hook_func(registry)
+                hooks_executed += 1
+                logger.debug(f"Executed custom commands hook: {hook_func.__name__}")
+            except Exception as e:
+                logger.error(f"Failed to execute custom commands hook {hook_func.__name__}: {e}")
+        
+        logger.info(f"Executed {hooks_executed} custom commands hooks")
+        return hooks_executed
     
-    def unregister_global_before_hook(self, hook: Callable[[HookContext], None]) -> None:
+    def execute_before_init_hooks(self) -> int:
         """
-        Unregister a global before hook.
+        Execute all registered before initialization hooks.
         
-        Args:
-            hook: Hook function to unregister
+        Returns:
+            Number of hooks executed.
         """
-        try:
-            self._global_before_hooks.remove(hook)
-            logger.debug("Unregistered global before hook")
-        except ValueError:
-            logger.warning("Global before hook not found")
+        logger.debug("Executing before init hooks...")
+        hooks_executed = 0
+        
+        for hook_func in self._before_init_hooks:
+            try:
+                hook_func()
+                hooks_executed += 1
+                logger.debug(f"Executed before init hook: {hook_func.__name__}")
+            except Exception as e:
+                logger.error(f"Failed to execute before init hook {hook_func.__name__}: {e}")
+        
+        logger.debug(f"Executed {hooks_executed} before init hooks")
+        return hooks_executed
     
-    def unregister_global_after_hook(self, hook: Callable[[HookContext], None]) -> None:
+    def execute_after_init_hooks(self) -> int:
         """
-        Unregister a global after hook.
+        Execute all registered after initialization hooks.
         
-        Args:
-            hook: Hook function to unregister
+        Returns:
+            Number of hooks executed.
         """
-        try:
-            self._global_after_hooks.remove(hook)
-            logger.debug("Unregistered global after hook")
-        except ValueError:
-            logger.warning("Global after hook not found")
+        logger.debug("Executing after init hooks...")
+        hooks_executed = 0
+        
+        for hook_func in self._after_init_hooks:
+            try:
+                hook_func()
+                hooks_executed += 1
+                logger.debug(f"Executed after init hook: {hook_func.__name__}")
+            except Exception as e:
+                logger.error(f"Failed to execute after init hook {hook_func.__name__}: {e}")
+        
+        logger.debug(f"Executed {hooks_executed} after init hooks")
+        return hooks_executed
     
-    def execute_before_hooks(self, command_name: str, params: Dict[str, Any]) -> HookContext:
+    def execute_before_command_hooks(self, command_name: str, params: Dict[str, Any]) -> int:
         """
-        Execute all before hooks for a command.
+        Execute all registered before command execution hooks.
         
         Args:
-            command_name: Name of the command
+            command_name: Name of the command being executed
             params: Command parameters
             
         Returns:
-            HookContext with execution results
+            Number of hooks executed.
         """
-        context = HookContext(
-            command_name=command_name,
-            params=params,
-            hook_type=HookType.BEFORE_EXECUTION
-        )
+        logger.debug(f"Executing before command hooks for: {command_name}")
+        hooks_executed = 0
         
-        # Execute global before hooks
-        for hook in self._global_before_hooks:
+        for hook_func in self._before_command_hooks:
             try:
-                hook(context)
+                hook_func(command_name, params)
+                hooks_executed += 1
+                logger.debug(f"Executed before command hook: {hook_func.__name__}")
             except Exception as e:
-                logger.error(f"Error in global before hook for command {command_name}: {e}")
+                logger.error(f"Failed to execute before command hook {hook_func.__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
+        logger.debug(f"Executed {hooks_executed} before command hooks")
+        return hooks_executed
     
-    def execute_after_hooks(self, command_name: str, params: Dict[str, Any], result: Any) -> HookContext:
+    def execute_after_command_hooks(self, command_name: str, params: Dict[str, Any], result: Any) -> int:
         """
-        Execute all after hooks for a command.
+        Execute all registered after command execution hooks.
         
         Args:
-            command_name: Name of the command
-            params: Command parameters
+            command_name: Name of the command that was executed
+            params: Command parameters that were used
             result: Command execution result
             
         Returns:
-            HookContext with execution results
+            Number of hooks executed.
         """
-        context = HookContext(
-            command_name=command_name,
-            params=params,
-            hook_type=HookType.AFTER_EXECUTION,
-            result=result
-        )
+        logger.debug(f"Executing after command hooks for: {command_name}")
+        hooks_executed = 0
         
-        # 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:
+        for hook_func in self._after_command_hooks:
             try:
-                hook(context)
+                hook_func(command_name, params, result)
+                hooks_executed += 1
+                logger.debug(f"Executed after command hook: {hook_func.__name__}")
             except Exception as e:
-                logger.error(f"Error in global after hook for command {command_name}: {e}")
+                logger.error(f"Failed to execute after command hook {hook_func.__name__}: {e}")
         
-        return context
+        logger.debug(f"Executed {hooks_executed} after command hooks")
+        return hooks_executed
     
-    def clear_hooks(self, command_name: Optional[str] = None) -> None:
+    def clear_hooks(self) -> 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.
+        Clear all registered 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)
-        }
+        self._custom_commands_hooks.clear()
+        self._before_init_hooks.clear()
+        self._after_init_hooks.clear()
+        self._before_command_hooks.clear()
+        self._after_command_hooks.clear()
+        logger.debug("Cleared all hooks")
 
 
 # Global hooks instance
-hooks = CommandHooks() 
+hooks = CommandHooks()
+
+
+def register_custom_commands_hook(hook_func: Callable) -> None:
+    """
+    Register a hook function for custom commands registration.
+    
+    Args:
+        hook_func: Function that registers custom commands.
+                  Should accept registry as parameter.
+    """
+    hooks.register_custom_commands_hook(hook_func)
+
+
+def register_before_init_hook(hook_func: Callable) -> None:
+    """
+    Register a hook function to be called before system initialization.
+    
+    Args:
+        hook_func: Function to call before initialization.
+    """
+    hooks.register_before_init_hook(hook_func)
+
+
+def register_after_init_hook(hook_func: Callable) -> None:
+    """
+    Register a hook function to be called after system initialization.
+    
+    Args:
+        hook_func: Function to call after initialization.
+    """
+    hooks.register_after_init_hook(hook_func)
+
+
+def register_before_command_hook(hook_func: Callable) -> None:
+    """
+    Register a hook function to be called before command execution.
+    
+    Args:
+        hook_func: Function to call before command execution.
+                  Should accept command_name and params as parameters.
+    """
+    hooks.register_before_command_hook(hook_func)
+
+
+def register_after_command_hook(hook_func: Callable) -> None:
+    """
+    Register a hook function to be called after command execution.
+    
+    Args:
+        hook_func: Function to call after command execution.
+                  Should accept command_name, params, and result as parameters.
+    """
+    hooks.register_after_command_hook(hook_func)