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

mcp_proxy_adapter/commands/proxy_registration_command.py

@@ -0,0 +1,409 @@
+"""
+Proxy Registration Command
+
+This command handles proxy registration functionality with security framework integration.
+It provides endpoints for registration, unregistration, heartbeat, and discovery.
+
+Author: Vasiliy Zdanovskiy
+email: vasilyvz@gmail.com
+"""
+
+import json
+import time
+import uuid
+from typing import Dict, Any, List, Optional
+from dataclasses import dataclass
+
+from mcp_proxy_adapter.commands.base import Command
+from mcp_proxy_adapter.commands.result import SuccessResult
+from mcp_proxy_adapter.core.logging import logger
+
+
+@dataclass
+class ProxyRegistrationCommandResult(SuccessResult):
+    """Result of proxy registration command."""
+    
+    operation: str
+    success: bool
+    server_key: Optional[str] = None
+    message: str = ""
+    details: Optional[Dict[str, Any]] = None
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert result to dictionary."""
+        result = {
+            "operation": self.operation,
+            "success": self.success,
+            "message": self.message
+        }
+        
+        if self.server_key:
+            result["server_key"] = self.server_key
+        
+        if self.details:
+            result["details"] = self.details
+        
+        return result
+    
+    @classmethod
+    def get_schema(cls) -> Dict[str, Any]:
+        """Get JSON schema for result."""
+        return {
+            "type": "object",
+            "properties": {
+                "operation": {
+                    "type": "string",
+                    "description": "Operation performed"
+                },
+                "success": {
+                    "type": "boolean",
+                    "description": "Whether operation was successful"
+                },
+                "server_key": {
+                    "type": "string",
+                    "description": "Server key for registered server"
+                },
+                "message": {
+                    "type": "string",
+                    "description": "Result message"
+                },
+                "details": {
+                    "type": "object",
+                    "description": "Additional details"
+                }
+            },
+            "required": ["operation", "success", "message"]
+        }
+
+
+class ProxyRegistrationCommand(Command):
+    """Proxy registration command with security framework integration."""
+    
+    name = "proxy_registration"
+    descr = "Proxy registration operations (register, unregister, heartbeat, discover)"
+    category = "proxy"
+    author = "Vasiliy Zdanovskiy"
+    email = "vasilyvz@gmail.com"
+    
+    # In-memory registry for testing
+    _registry: Dict[str, Dict[str, Any]] = {}
+    _server_counter = 1
+    
+    async def execute(self, **kwargs) -> ProxyRegistrationCommandResult:
+        """
+        Execute proxy registration command.
+        
+        Args:
+            operation: Operation to perform (register, unregister, heartbeat, discover)
+            server_id: Server ID for registration
+            server_url: Server URL for registration
+            server_name: Server name
+            description: Server description
+            version: Server version
+            capabilities: Server capabilities
+            endpoints: Server endpoints
+            auth_method: Authentication method
+            security_enabled: Whether security is enabled
+            server_key: Server key for unregistration/heartbeat
+            copy_number: Copy number for unregistration
+            timestamp: Timestamp for heartbeat
+            status: Status for heartbeat
+            
+        Returns:
+            ProxyRegistrationCommandResult
+        """
+        operation = kwargs.get("operation", "register")
+        
+        # Check user permissions
+        context = kwargs.get("context", {})
+        user_info = context.get("user", {})
+        user_permissions = user_info.get("permissions", [])
+        
+        # Define required permissions for each operation
+        operation_permissions = {
+            "register": ["register"],
+            "unregister": ["unregister"],
+            "heartbeat": ["heartbeat"],
+            "discover": ["discover"]
+        }
+        
+        required_permissions = operation_permissions.get(operation, ["read"])
+        
+        # Check if user has required permissions
+        logger.info(f"Checking permissions: user_permissions={user_permissions}, required={required_permissions}")
+        if not self._check_permissions(user_permissions, required_permissions):
+            return ProxyRegistrationCommandResult(
+                operation=operation,
+                success=False,
+                message=f"Permission denied: {operation} requires {required_permissions}"
+            )
+        
+        logger.info(f"Executing proxy registration operation: {operation}")
+        logger.debug(f"User permissions: {user_permissions}, required: {required_permissions}")
+        
+        if operation == "register":
+            return await self._handle_register(kwargs)
+        elif operation == "unregister":
+            return await self._handle_unregister(kwargs)
+        elif operation == "heartbeat":
+            return await self._handle_heartbeat(kwargs)
+        elif operation == "discover":
+            return await self._handle_discover(kwargs)
+        else:
+            return ProxyRegistrationCommandResult(
+                operation=operation,
+                success=False,
+                message=f"Unknown operation: {operation}"
+            )
+    
+    async def _handle_register(self, kwargs: Dict[str, Any]) -> ProxyRegistrationCommandResult:
+        """Handle registration operation."""
+        server_id = kwargs.get("server_id")
+        server_url = kwargs.get("server_url")
+        server_name = kwargs.get("server_name", "Unknown Server")
+        description = kwargs.get("description", "")
+        version = kwargs.get("version", "1.0.0")
+        capabilities = kwargs.get("capabilities", ["jsonrpc", "rest"])
+        endpoints = kwargs.get("endpoints", {})
+        auth_method = kwargs.get("auth_method", "none")
+        security_enabled = kwargs.get("security_enabled", False)
+        
+        if not server_id or not server_url:
+            return ProxyRegistrationCommandResult(
+                operation="register",
+                success=False,
+                message="Missing required parameters: server_id and server_url"
+            )
+        
+        # Check if server already exists
+        existing_servers = [key for key in self._registry.keys() if key.startswith(server_id)]
+        copy_number = len(existing_servers) + 1
+        server_key = f"{server_id}_{copy_number}"
+        
+        # Create server record
+        server_record = {
+            "server_id": server_id,
+            "server_url": server_url,
+            "server_name": server_name,
+            "description": description,
+            "version": version,
+            "capabilities": capabilities,
+            "endpoints": endpoints,
+            "auth_method": auth_method,
+            "security_enabled": security_enabled,
+            "registered_at": int(time.time()),
+            "last_heartbeat": int(time.time()),
+            "status": "active"
+        }
+        
+        self._registry[server_key] = server_record
+        
+        logger.info(f"Registered server: {server_key} at {server_url}")
+        
+        return ProxyRegistrationCommandResult(
+            operation="register",
+            success=True,
+            server_key=server_key,
+            message=f"Server registered successfully with key: {server_key}",
+            details={
+                "server_id": server_id,
+                "copy_number": copy_number,
+                "registered_at": server_record["registered_at"]
+            }
+        )
+    
+    async def _handle_unregister(self, kwargs: Dict[str, Any]) -> ProxyRegistrationCommandResult:
+        """Handle unregistration operation."""
+        server_id = kwargs.get("server_id")
+        copy_number = kwargs.get("copy_number", 1)
+        
+        if not server_id:
+            return ProxyRegistrationCommandResult(
+                operation="unregister",
+                success=False,
+                message="Missing required parameter: server_id"
+            )
+        
+        server_key = f"{server_id}_{copy_number}"
+        
+        if server_key in self._registry:
+            del self._registry[server_key]
+            logger.info(f"Unregistered server: {server_key}")
+            
+            return ProxyRegistrationCommandResult(
+                operation="unregister",
+                success=True,
+                message=f"Server unregistered successfully: {server_key}",
+                details={"unregistered": True}
+            )
+        else:
+            return ProxyRegistrationCommandResult(
+                operation="unregister",
+                success=True,
+                message=f"Server not found in registry: {server_key}",
+                details={"unregistered": False}
+            )
+    
+    async def _handle_heartbeat(self, kwargs: Dict[str, Any]) -> ProxyRegistrationCommandResult:
+        """Handle heartbeat operation."""
+        server_id = kwargs.get("server_id")
+        server_key = kwargs.get("server_key")
+        timestamp = kwargs.get("timestamp", int(time.time()))
+        status = kwargs.get("status", "healthy")
+        
+        if not server_key:
+            return ProxyRegistrationCommandResult(
+                operation="heartbeat",
+                success=False,
+                message="Missing required parameter: server_key"
+            )
+        
+        if server_key in self._registry:
+            self._registry[server_key]["last_heartbeat"] = timestamp
+            self._registry[server_key]["status"] = status
+            
+            logger.debug(f"Heartbeat received for server: {server_key}")
+            
+            return ProxyRegistrationCommandResult(
+                operation="heartbeat",
+                success=True,
+                message="Heartbeat processed successfully",
+                details={
+                    "server_key": server_key,
+                    "timestamp": timestamp,
+                    "status": status
+                }
+            )
+        else:
+            return ProxyRegistrationCommandResult(
+                operation="heartbeat",
+                success=False,
+                message=f"Server not found: {server_key}"
+            )
+    
+    async def _handle_discover(self, kwargs: Dict[str, Any]) -> ProxyRegistrationCommandResult:
+        """Handle discovery operation."""
+        # Return all registered servers
+        proxies = []
+        
+        for server_key, server_record in self._registry.items():
+            # Check if server is active (heartbeat within last 5 minutes)
+            last_heartbeat = server_record.get("last_heartbeat", 0)
+            if time.time() - last_heartbeat < 300:  # 5 minutes
+                proxy_info = {
+                    "server_key": server_key,
+                    "server_id": server_record["server_id"],
+                    "server_url": server_record["server_url"],
+                    "server_name": server_record["server_name"],
+                    "description": server_record["description"],
+                    "version": server_record["version"],
+                    "capabilities": server_record["capabilities"],
+                    "endpoints": server_record["endpoints"],
+                    "auth_method": server_record["auth_method"],
+                    "security_enabled": server_record["security_enabled"],
+                    "registered_at": server_record["registered_at"],
+                    "last_heartbeat": server_record["last_heartbeat"],
+                    "status": server_record["status"]
+                }
+                proxies.append(proxy_info)
+        
+        logger.info(f"Discovery request returned {len(proxies)} active servers")
+        
+        return ProxyRegistrationCommandResult(
+            operation="discover",
+            success=True,
+            message=f"Found {len(proxies)} active proxy servers",
+            details={"proxies": proxies}
+        )
+    
+    def _check_permissions(self, user_permissions: List[str], required_permissions: List[str]) -> bool:
+        """
+        Check if user has required permissions.
+        
+        Args:
+            user_permissions: User's permissions
+            required_permissions: Required permissions
+            
+        Returns:
+            True if user has required permissions
+        """
+        # Admin has all permissions
+        if "*" in user_permissions:
+            return True
+        
+        # Check if user has all required permissions
+        for required in required_permissions:
+            if required not in user_permissions:
+                return False
+        
+        return True
+    
+    @classmethod
+    def get_schema(cls) -> Dict[str, Any]:
+        """Get JSON schema for command parameters."""
+        return {
+            "type": "object",
+            "properties": {
+                "operation": {
+                    "type": "string",
+                    "enum": ["register", "unregister", "heartbeat", "discover"],
+                    "description": "Operation to perform",
+                    "default": "register"
+                },
+                "server_id": {
+                    "type": "string",
+                    "description": "Server ID for registration"
+                },
+                "server_url": {
+                    "type": "string",
+                    "description": "Server URL for registration"
+                },
+                "server_name": {
+                    "type": "string",
+                    "description": "Server name"
+                },
+                "description": {
+                    "type": "string",
+                    "description": "Server description"
+                },
+                "version": {
+                    "type": "string",
+                    "description": "Server version"
+                },
+                "capabilities": {
+                    "type": "array",
+                    "items": {"type": "string"},
+                    "description": "Server capabilities"
+                },
+                "endpoints": {
+                    "type": "object",
+                    "description": "Server endpoints"
+                },
+                "auth_method": {
+                    "type": "string",
+                    "description": "Authentication method"
+                },
+                "security_enabled": {
+                    "type": "boolean",
+                    "description": "Whether security is enabled"
+                },
+                "server_key": {
+                    "type": "string",
+                    "description": "Server key for unregistration/heartbeat"
+                },
+                "copy_number": {
+                    "type": "integer",
+                    "description": "Copy number for unregistration"
+                },
+                "timestamp": {
+                    "type": "integer",
+                    "description": "Timestamp for heartbeat"
+                },
+                "status": {
+                    "type": "string",
+                    "description": "Status for heartbeat"
+                }
+            },
+            "required": ["operation"],
+            "additionalProperties": False
+        } 

mcp_proxy_adapter/commands/reload_command.py

@@ -20,11 +20,11 @@ class ReloadResult:
     def __init__(
         self,
         config_reloaded: bool,
-        commands_discovered: int,
-        custom_commands_preserved: int,
-        total_commands: int,
-        built_in_commands: int,
+        builtin_commands: int,
         custom_commands: int,
+        loaded_commands: int,
+        remote_commands: int = 0,
+        total_commands: int = 0,
         server_restart_required: bool = True,
         success: bool = True,
         error_message: Optional[str] = None
@@ -34,21 +34,20 @@ class ReloadResult:
         
         Args:
             config_reloaded: Whether configuration was reloaded successfully
-            commands_discovered: Number of commands discovered
-            custom_commands_preserved: Number of custom commands preserved
+            builtin_commands: Number of built-in commands registered
+            custom_commands: Number of custom commands registered
+            loaded_commands: Number of commands loaded from directory
             total_commands: Total number of commands after reload
-            built_in_commands: Number of built-in commands
-            custom_commands: Number of custom commands
             server_restart_required: Whether server restart is required
             success: Whether reload was successful
             error_message: Error message if reload failed
         """
         self.config_reloaded = config_reloaded
-        self.commands_discovered = commands_discovered
-        self.custom_commands_preserved = custom_commands_preserved
-        self.total_commands = total_commands
-        self.built_in_commands = built_in_commands
+        self.builtin_commands = builtin_commands
         self.custom_commands = custom_commands
+        self.loaded_commands = loaded_commands
+        self.remote_commands = remote_commands
+        self.total_commands = total_commands
         self.server_restart_required = server_restart_required
         self.success = success
         self.error_message = error_message
@@ -63,11 +62,11 @@ class ReloadResult:
         return {
             "success": self.success,
             "config_reloaded": self.config_reloaded,
-            "commands_discovered": self.commands_discovered,
-            "custom_commands_preserved": self.custom_commands_preserved,
-            "total_commands": self.total_commands,
-            "built_in_commands": self.built_in_commands,
+            "builtin_commands": self.builtin_commands,
             "custom_commands": self.custom_commands,
+            "loaded_commands": self.loaded_commands,
+            "remote_commands": self.remote_commands,
+            "total_commands": self.total_commands,
             "server_restart_required": self.server_restart_required,
             "message": "Server restart required to apply configuration changes",
             "error_message": self.error_message
@@ -91,25 +90,25 @@ class ReloadResult:
                     "type": "boolean",
                     "description": "Whether configuration was reloaded successfully"
                 },
-                "commands_discovered": {
+                "builtin_commands": {
                     "type": "integer",
-                    "description": "Number of commands discovered"
+                    "description": "Number of built-in commands registered"
                 },
-                "custom_commands_preserved": {
+                "custom_commands": {
                     "type": "integer",
-                    "description": "Number of custom commands preserved"
+                    "description": "Number of custom commands registered"
                 },
-                "total_commands": {
+                "loaded_commands": {
                     "type": "integer",
-                    "description": "Total number of commands after reload"
+                    "description": "Number of commands loaded from directory"
                 },
-                "built_in_commands": {
+                "remote_commands": {
                     "type": "integer",
-                    "description": "Number of built-in commands"
+                    "description": "Number of commands loaded from remote plugins"
                 },
-                "custom_commands": {
+                "total_commands": {
                     "type": "integer",
-                    "description": "Number of custom commands"
+                    "description": "Total number of commands after reload"
                 },
                 "server_restart_required": {
                     "type": "boolean",
@@ -125,9 +124,8 @@ class ReloadResult:
                 }
             },
             "required": [
-                "success", "config_reloaded", "commands_discovered",
-                "custom_commands_preserved", "total_commands",
-                "built_in_commands", "custom_commands", "server_restart_required"
+                "success", "config_reloaded", "builtin_commands", "custom_commands",
+                "loaded_commands", "remote_commands", "total_commands", "server_restart_required"
             ]
         }
 
@@ -135,7 +133,7 @@ class ReloadResult:
 class ReloadCommand(Command):
     """
     Command for reloading configuration and rediscovering commands.
-    Note: This command will trigger a server restart to apply configuration changes.
+    Uses the unified initialization logic.
     """
     
     name = "reload"
@@ -145,7 +143,7 @@ class ReloadCommand(Command):
         Execute reload command.
         
         Args:
-            **params: Command parameters (currently unused)
+            **params: Command parameters (config_path)
             
         Returns:
             ReloadResult with reload information
@@ -153,18 +151,23 @@ class ReloadCommand(Command):
         try:
             logger.info("🔄 Starting configuration and commands reload...")
             
-            # Perform reload
-            reload_info = registry.reload_config_and_commands()
+            # Get config path from parameters
+            config_path = params.get("config_path")
+            if not config_path:
+                logger.warning("No config_path provided, using default configuration")
+            
+            # Perform reload using unified initialization
+            reload_info = registry.reload_system(config_path=config_path)
             
             # Create result
             result = ReloadResult(
                 config_reloaded=reload_info.get("config_reloaded", False),
-                commands_discovered=reload_info.get("commands_discovered", 0),
-                custom_commands_preserved=reload_info.get("custom_commands_preserved", 0),
-                total_commands=reload_info.get("total_commands", 0),
-                built_in_commands=reload_info.get("built_in_commands", 0),
+                builtin_commands=reload_info.get("builtin_commands", 0),
                 custom_commands=reload_info.get("custom_commands", 0),
-                server_restart_required=True,
+                loaded_commands=reload_info.get("loaded_commands", 0),
+                remote_commands=reload_info.get("remote_commands", 0),
+                total_commands=reload_info.get("total_commands", 0),
+                server_restart_required=True,  # Default to True as per tests
                 success=True
             )
             
@@ -175,11 +178,11 @@ class ReloadCommand(Command):
             logger.error(f"❌ Reload failed: {str(e)}")
             return ReloadResult(
                 config_reloaded=False,
-                commands_discovered=0,
-                custom_commands_preserved=0,
-                total_commands=0,
-                built_in_commands=0,
+                builtin_commands=0,
                 custom_commands=0,
+                loaded_commands=0,
+                remote_commands=0,
+                total_commands=0,
                 server_restart_required=False,
                 success=False,
                 error_message=str(e)
@@ -196,15 +199,10 @@ class ReloadCommand(Command):
         return {
             "type": "object",
             "properties": {
-                "package_path": {
+                "config_path": {
                     "type": "string",
-                    "description": "Path to package with commands to discover",
-                    "default": "mcp_proxy_adapter.commands"
-                },
-                "force_restart": {
-                    "type": "boolean",
-                    "description": "Force server restart to apply configuration changes",
-                    "default": True
+                    "description": "Path to configuration file to reload",
+                    "default": None
                 }
             },
             "additionalProperties": False

mcp_proxy_adapter/commands/result.py

@@ -150,6 +150,7 @@ class ErrorResult(CommandResult):
             details: Additional error details.
         """
         self.message = message
+        self.error = message  # For backward compatibility with tests
         self.code = code
         self.details = details or {}