mcp-proxy-adapter 6.6.9__py3-none-any.whl → 6.7.1__py3-none-any.whl

mcp_proxy_adapter/api/app.py

@@ -159,18 +159,24 @@ def create_lifespan(config_path: Optional[str] = None):
         # Determine registration URL using unified logic
         early_server_url = _determine_registration_url(config)
         try:
-            from mcp_proxy_adapter.core.proxy_registration import (
-                proxy_registration_manager,
+            from mcp_proxy_adapter.core.async_proxy_registration import (
+                initialize_async_registration,
+                start_async_registration,
             )
 
-            if proxy_registration_manager is not None:
-                proxy_registration_manager.set_server_url(early_server_url)
-                logger.info(
-                    "🔍 Early injection of server_url for registration: %s",
-                    early_server_url,
-                )
+            # Initialize async registration manager
+            async_manager = initialize_async_registration(config)
+            logger.info(
+                "🔍 Initialized async registration manager with server_url: %s",
+                early_server_url,
+            )
+            
+            # Start async registration in background thread
+            start_async_registration(early_server_url)
+            logger.info("🚀 Started async proxy registration and heartbeat thread")
+            
         except Exception as e:
-            logger.error(f"Failed to inject early server_url: {e}")
+            logger.error(f"Failed to initialize async registration: {e}")
 
         # Initialize system using unified logic (may perform registration)
         if config_path:
@@ -194,6 +200,27 @@ def create_lifespan(config_path: Optional[str] = None):
 
             # Determine registration URL using unified logic
             server_url = _determine_registration_url(final_config)
+            
+            # Update async registration manager with final server URL
+            try:
+                from mcp_proxy_adapter.core.async_proxy_registration import (
+                    get_async_registration_manager,
+                )
+                
+                async_manager = get_async_registration_manager()
+                if async_manager:
+                    async_manager.set_server_url(server_url)
+                    logger.info(f"🔍 Updated async registration manager with final server_url: {server_url}")
+                    
+                    # Get current status
+                    status = async_manager.get_status()
+                    logger.info(f"📊 Registration status: {status}")
+                else:
+                    logger.warning("Async registration manager not available")
+                    
+            except Exception as e:
+                logger.error(f"Failed to update async registration manager: {e}")
+            
             try:
                 print("🔍 Registration server_url resolved to (print):", server_url)
             except Exception:
@@ -228,12 +255,30 @@ def create_lifespan(config_path: Optional[str] = None):
         # Shutdown events
         logger.info("Application shutting down")
 
-        # Unregister from proxy if enabled
-        unregistration_success = await unregister_from_proxy()
-        if unregistration_success:
-            logger.info("✅ Proxy unregistration completed successfully")
-        else:
-            logger.warning("⚠️ Proxy unregistration failed or was disabled")
+        # Stop async registration manager (this will also unregister)
+        try:
+            from mcp_proxy_adapter.core.async_proxy_registration import (
+                stop_async_registration,
+                get_registration_status,
+            )
+            
+            # Get final status before stopping
+            final_status = get_registration_status()
+            logger.info(f"📊 Final registration status: {final_status}")
+            
+            # Stop async registration (this will also unregister)
+            stop_async_registration()
+            logger.info("✅ Async registration manager stopped successfully")
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to stop async registration: {e}")
+            
+            # Fallback to old unregistration method
+            unregistration_success = await unregister_from_proxy()
+            if unregistration_success:
+                logger.info("✅ Fallback proxy unregistration completed successfully")
+            else:
+                logger.warning("⚠️ Fallback proxy unregistration failed or was disabled")
 
     return lifespan
 

mcp_proxy_adapter/commands/__init__.py

@@ -24,6 +24,9 @@ from mcp_proxy_adapter.commands.echo_command import EchoCommand
 from mcp_proxy_adapter.commands.proxy_registration_command import (
     ProxyRegistrationCommand,
 )
+from mcp_proxy_adapter.commands.registration_status_command import (
+    RegistrationStatusCommand,
+)
 
 __all__ = [
     "Command",
@@ -43,4 +46,5 @@ __all__ = [
     "RoleTestCommand",
     "EchoCommand",
     "ProxyRegistrationCommand",
+    "RegistrationStatusCommand",
 ]

mcp_proxy_adapter/commands/registration_status_command.py

@@ -0,0 +1,119 @@
+"""
+Registration Status Command
+
+This command provides information about the current proxy registration status,
+including async registration state, heartbeat status, and statistics.
+
+Author: Vasiliy Zdanovskiy
+email: vasilyvz@gmail.com
+"""
+
+from typing import Dict, Any
+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 RegistrationStatusCommandResult(SuccessResult):
+    """Result of registration status command."""
+
+    status: Dict[str, Any]
+    message: str = "Registration status retrieved successfully"
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert result to dictionary."""
+        return {
+            "success": True,
+            "status": self.status,
+            "message": self.message,
+        }
+
+    @classmethod
+    def get_schema(cls) -> Dict[str, Any]:
+        """Get JSON schema for result."""
+        return {
+            "type": "object",
+            "properties": {
+                "success": {
+                    "type": "boolean",
+                    "description": "Whether the command was successful",
+                },
+                "status": {
+                    "type": "object",
+                    "description": "Registration status information",
+                    "properties": {
+                        "state": {"type": "string", "description": "Current registration state"},
+                        "server_key": {"type": "string", "description": "Server key if registered"},
+                        "last_attempt": {"type": "number", "description": "Timestamp of last registration attempt"},
+                        "last_success": {"type": "number", "description": "Timestamp of last successful registration"},
+                        "last_error": {"type": "string", "description": "Last error message"},
+                        "attempt_count": {"type": "integer", "description": "Total registration attempts"},
+                        "success_count": {"type": "integer", "description": "Total successful registrations"},
+                        "heartbeat_enabled": {"type": "boolean", "description": "Whether heartbeat is enabled"},
+                        "heartbeat_interval": {"type": "integer", "description": "Heartbeat interval in seconds"},
+                        "thread_alive": {"type": "boolean", "description": "Whether registration thread is alive"},
+                    },
+                },
+                "message": {"type": "string", "description": "Result message"},
+            },
+            "required": ["success", "status", "message"],
+        }
+
+
+class RegistrationStatusCommand(Command):
+    """Command to get proxy registration status."""
+
+    name = "registration_status"
+    descr = "Get current proxy registration status and statistics"
+    category = "proxy"
+    author = "Vasiliy Zdanovskiy"
+    email = "vasilyvz@gmail.com"
+
+    async def execute(self, **kwargs) -> RegistrationStatusCommandResult:
+        """
+        Execute registration status command.
+
+        Returns:
+            RegistrationStatusCommandResult with current status
+        """
+        logger.info("Executing registration status command")
+
+        try:
+            from mcp_proxy_adapter.core.async_proxy_registration import (
+                get_registration_status,
+            )
+
+            status = get_registration_status()
+            
+            logger.info(f"Registration status retrieved: {status}")
+
+            return RegistrationStatusCommandResult(
+                status=status,
+                message="Registration status retrieved successfully"
+            )
+
+        except Exception as e:
+            logger.error(f"Failed to get registration status: {e}")
+            
+            error_status = {
+                "state": "error",
+                "error": str(e),
+                "thread_alive": False
+            }
+            
+            return RegistrationStatusCommandResult(
+                status=error_status,
+                message=f"Failed to get registration status: {e}"
+            )
+
+    @classmethod
+    def get_schema(cls) -> Dict[str, Any]:
+        """Get JSON schema for command parameters."""
+        return {
+            "type": "object",
+            "properties": {},
+            "additionalProperties": False,
+        }

mcp_proxy_adapter/core/app_runner.py

@@ -16,6 +16,7 @@ from typing import Dict, Any, List, Optional
 from fastapi import FastAPI
 
 from mcp_proxy_adapter.core.logging import get_logger
+from mcp_proxy_adapter.core.signal_handler import setup_signal_handling, is_shutdown_requested
 
 logger = get_logger("app_runner")
 
@@ -250,6 +251,30 @@ class ApplicationRunner:
                 print(f"  - {error}", file=sys.stderr)
             sys.exit(1)
 
+        # Setup signal handling for graceful shutdown
+        def shutdown_callback():
+            """Callback for graceful shutdown with proxy unregistration."""
+            print("\n🛑 Graceful shutdown initiated...")
+            try:
+                from mcp_proxy_adapter.core.async_proxy_registration import (
+                    stop_async_registration,
+                    get_registration_status,
+                )
+                
+                # Get final status
+                final_status = get_registration_status()
+                print(f"📊 Final registration status: {final_status}")
+                
+                # Stop async registration (this will unregister from proxy)
+                stop_async_registration()
+                print("✅ Proxy unregistration completed")
+                
+            except Exception as e:
+                print(f"❌ Error during shutdown: {e}")
+        
+        setup_signal_handling(shutdown_callback)
+        print("🔧 Signal handling configured for graceful shutdown")
+
         # Setup hooks
         self.setup_hooks()
 
@@ -276,14 +301,16 @@ class ApplicationRunner:
             import asyncio
 
             print(f"🚀 Starting server on {host}:{port}")
-            print("   Use Ctrl+C to stop the server")
+            print("🛑 Use Ctrl+C or send SIGTERM for graceful shutdown")
             print("=" * 60)
 
             # Run with hypercorn
             asyncio.run(hypercorn.asyncio.serve(self.app, **server_kwargs))
 
         except KeyboardInterrupt:
-            print("\n🛑 Server stopped by user")
+            print("\n🛑 Server stopped by user (Ctrl+C)")
+            if is_shutdown_requested():
+                print("✅ Graceful shutdown completed")
         except Exception as e:
             print(f"\n❌ Failed to start server: {e}", file=sys.stderr)
             sys.exit(1)

mcp_proxy_adapter/core/async_proxy_registration.py

@@ -0,0 +1,285 @@
+"""
+Asynchronous proxy registration with heartbeat functionality.
+
+This module provides automatic proxy registration in a separate thread
+with continuous heartbeat monitoring and automatic re-registration on failures.
+
+Author: Vasiliy Zdanovskiy
+email: vasilyvz@gmail.com
+"""
+
+import asyncio
+import time
+import threading
+from typing import Dict, Any, Optional
+from dataclasses import dataclass
+from enum import Enum
+
+from mcp_proxy_adapter.core.logging import logger
+from mcp_proxy_adapter.core.proxy_registration import ProxyRegistrationManager
+
+
+class RegistrationState(Enum):
+    """Registration state enumeration."""
+    DISABLED = "disabled"
+    NOT_REGISTERED = "not_registered"
+    REGISTERING = "registering"
+    REGISTERED = "registered"
+    HEARTBEAT_FAILED = "heartbeat_failed"
+    ERROR = "error"
+
+
+@dataclass
+class RegistrationStatus:
+    """Registration status information."""
+    state: RegistrationState
+    server_key: Optional[str] = None
+    last_attempt: Optional[float] = None
+    last_success: Optional[float] = None
+    last_error: Optional[str] = None
+    attempt_count: int = 0
+    success_count: int = 0
+
+
+class AsyncProxyRegistrationManager:
+    """
+    Asynchronous proxy registration manager with heartbeat functionality.
+    
+    Runs registration and heartbeat in a separate thread, automatically
+    handling re-registration on failures and continuous monitoring.
+    """
+    
+    def __init__(self, config: Dict[str, Any]):
+        """
+        Initialize the async proxy registration manager.
+        
+        Args:
+            config: Application configuration
+        """
+        self.config = config
+        self.registration_manager = ProxyRegistrationManager(config)
+        self.status = RegistrationStatus(state=RegistrationState.NOT_REGISTERED)
+        self._thread: Optional[threading.Thread] = None
+        self._stop_event = threading.Event()
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+        
+        # Get heartbeat configuration
+        heartbeat_config = config.get("registration", {}).get("heartbeat", {})
+        self.heartbeat_enabled = heartbeat_config.get("enabled", True)
+        self.heartbeat_interval = heartbeat_config.get("interval", 30)
+        self.heartbeat_timeout = heartbeat_config.get("timeout", 10)
+        self.retry_attempts = heartbeat_config.get("retry_attempts", 3)
+        self.retry_delay = heartbeat_config.get("retry_delay", 5)
+        
+        logger.info(f"AsyncProxyRegistrationManager initialized: heartbeat_enabled={self.heartbeat_enabled}, interval={self.heartbeat_interval}s")
+    
+    def is_enabled(self) -> bool:
+        """Check if proxy registration is enabled."""
+        return self.registration_manager.is_enabled()
+    
+    def set_server_url(self, server_url: str):
+        """Set the server URL for registration."""
+        self.registration_manager.set_server_url(server_url)
+        logger.info(f"Server URL set for async registration: {server_url}")
+    
+    def start(self):
+        """Start the async registration and heartbeat thread."""
+        if not self.is_enabled():
+            logger.info("Proxy registration is disabled, not starting async manager")
+            self.status.state = RegistrationState.DISABLED
+            return
+        
+        if self._thread and self._thread.is_alive():
+            logger.warning("Async registration thread is already running")
+            return
+        
+        logger.info("Starting async proxy registration and heartbeat thread")
+        self._stop_event.clear()
+        self._thread = threading.Thread(target=self._run_async_loop, daemon=True)
+        self._thread.start()
+    
+    def stop(self):
+        """Stop the async registration and heartbeat thread."""
+        if not self._thread or not self._thread.is_alive():
+            return
+        
+        logger.info("Stopping async proxy registration and heartbeat thread")
+        self._stop_event.set()
+        
+        # Unregister if registered
+        if self.status.state == RegistrationState.REGISTERED:
+            try:
+                # Run unregistration in the thread's event loop
+                if self._loop and not self._loop.is_closed():
+                    future = asyncio.run_coroutine_threadsafe(
+                        self.registration_manager.unregister_server(), 
+                        self._loop
+                    )
+                    future.result(timeout=10)
+                    logger.info("Successfully unregistered from proxy during shutdown")
+            except Exception as e:
+                logger.error(f"Failed to unregister during shutdown: {e}")
+        
+        self._thread.join(timeout=5)
+        if self._thread.is_alive():
+            logger.warning("Async registration thread did not stop gracefully")
+    
+    def _run_async_loop(self):
+        """Run the async event loop in the thread."""
+        try:
+            self._loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(self._loop)
+            self._loop.run_until_complete(self._async_main())
+        except Exception as e:
+            logger.error(f"Async registration loop error: {e}")
+        finally:
+            if self._loop and not self._loop.is_closed():
+                self._loop.close()
+    
+    async def _async_main(self):
+        """Main async loop for registration and heartbeat."""
+        logger.info("Async registration loop started")
+        
+        # Initial registration attempt
+        await self._attempt_registration()
+        
+        # Main loop
+        while not self._stop_event.is_set():
+            try:
+                if self.status.state == RegistrationState.REGISTERED:
+                    if self.heartbeat_enabled:
+                        # Perform heartbeat
+                        await self._perform_heartbeat()
+                    else:
+                        # Just wait if heartbeat is disabled
+                        await asyncio.sleep(self.heartbeat_interval)
+                else:
+                    # Try to register
+                    await self._attempt_registration()
+                    # Wait before retry
+                    await asyncio.sleep(self.retry_delay)
+                    
+            except Exception as e:
+                logger.error(f"Error in async registration loop: {e}")
+                self.status.state = RegistrationState.ERROR
+                self.status.last_error = str(e)
+                await asyncio.sleep(self.retry_delay)
+        
+        logger.info("Async registration loop stopped")
+    
+    async def _attempt_registration(self):
+        """Attempt to register with the proxy."""
+        if self.status.state == RegistrationState.REGISTERING:
+            return  # Already registering
+        
+        self.status.state = RegistrationState.REGISTERING
+        self.status.attempt_count += 1
+        self.status.last_attempt = time.time()
+        
+        logger.info(f"Attempting proxy registration (attempt #{self.status.attempt_count})")
+        
+        try:
+            success = await self.registration_manager.register_server()
+            
+            if success:
+                self.status.state = RegistrationState.REGISTERED
+                self.status.last_success = time.time()
+                self.status.success_count += 1
+                self.status.last_error = None
+                logger.info(f"✅ Proxy registration successful (attempt #{self.status.attempt_count})")
+            else:
+                self.status.state = RegistrationState.NOT_REGISTERED
+                self.status.last_error = "Registration returned False"
+                logger.warning(f"❌ Proxy registration failed (attempt #{self.status.attempt_count}): Registration returned False")
+                
+        except Exception as e:
+            self.status.state = RegistrationState.NOT_REGISTERED
+            self.status.last_error = str(e)
+            logger.error(f"❌ Proxy registration failed (attempt #{self.status.attempt_count}): {e}")
+    
+    async def _perform_heartbeat(self):
+        """Perform heartbeat to the proxy."""
+        try:
+            # Use the registration manager's heartbeat functionality
+            success = await self.registration_manager.heartbeat()
+            
+            if success:
+                logger.debug("💓 Heartbeat successful")
+                self.status.last_success = time.time()
+            else:
+                logger.warning("💓 Heartbeat failed")
+                self.status.state = RegistrationState.HEARTBEAT_FAILED
+                self.status.last_error = "Heartbeat failed"
+                
+                # Try to re-register after heartbeat failure
+                logger.info("Attempting re-registration after heartbeat failure")
+                await self._attempt_registration()
+                
+        except Exception as e:
+            logger.error(f"💓 Heartbeat error: {e}")
+            self.status.state = RegistrationState.HEARTBEAT_FAILED
+            self.status.last_error = str(e)
+            
+            # Try to re-register after heartbeat error
+            logger.info("Attempting re-registration after heartbeat error")
+            await self._attempt_registration()
+        
+        # Wait for next heartbeat
+        await asyncio.sleep(self.heartbeat_interval)
+    
+    def get_status(self) -> Dict[str, Any]:
+        """Get current registration status."""
+        return {
+            "state": self.status.state.value,
+            "server_key": self.status.server_key,
+            "last_attempt": self.status.last_attempt,
+            "last_success": self.status.last_success,
+            "last_error": self.status.last_error,
+            "attempt_count": self.status.attempt_count,
+            "success_count": self.status.success_count,
+            "heartbeat_enabled": self.heartbeat_enabled,
+            "heartbeat_interval": self.heartbeat_interval,
+            "thread_alive": self._thread.is_alive() if self._thread else False
+        }
+
+
+# Global instance
+_async_registration_manager: Optional[AsyncProxyRegistrationManager] = None
+
+
+def get_async_registration_manager() -> Optional[AsyncProxyRegistrationManager]:
+    """Get the global async registration manager instance."""
+    return _async_registration_manager
+
+
+def initialize_async_registration(config: Dict[str, Any]) -> AsyncProxyRegistrationManager:
+    """Initialize the global async registration manager."""
+    global _async_registration_manager
+    _async_registration_manager = AsyncProxyRegistrationManager(config)
+    return _async_registration_manager
+
+
+def start_async_registration(server_url: str):
+    """Start async registration with the given server URL."""
+    global _async_registration_manager
+    if _async_registration_manager:
+        _async_registration_manager.set_server_url(server_url)
+        _async_registration_manager.start()
+    else:
+        logger.error("Async registration manager not initialized")
+
+
+def stop_async_registration():
+    """Stop async registration."""
+    global _async_registration_manager
+    if _async_registration_manager:
+        _async_registration_manager.stop()
+
+
+def get_registration_status() -> Dict[str, Any]:
+    """Get current registration status."""
+    global _async_registration_manager
+    if _async_registration_manager:
+        return _async_registration_manager.get_status()
+    else:
+        return {"state": "not_initialized"}