mcp-mesh 0.7.21__py3-none-any.whl → 0.8.0b1__py3-none-any.whl

_mcp_mesh/shared/server_discovery.py

@@ -8,8 +8,8 @@ outside the pipeline, such as by immediate uvicorn start in decorators.
 import gc
 import logging
 import socket
-from typing import Any, Dict, List, Optional, Tuple
 import threading
+from typing import Any, Dict, List, Optional, Tuple
 
 logger = logging.getLogger(__name__)
 
@@ -21,10 +21,10 @@ class ServerDiscoveryUtil:
     def discover_fastapi_instances() -> Dict[str, Dict[str, Any]]:
         """
         Discover FastAPI application instances in the Python runtime.
-        
+
         Uses intelligent deduplication to handle standard uvicorn patterns where
         the same app might be imported multiple times (e.g., "module:app" pattern).
-        
+
         Returns:
             Dict mapping app_id -> app_info where app_info contains:
             - 'instance': The FastAPI app instance
@@ -34,7 +34,7 @@ class ServerDiscoveryUtil:
         """
         fastapi_apps = {}
         seen_apps = {}  # For deduplication: title -> app_info
-        
+
         try:
             # Import FastAPI here to avoid dependency if not used
             from fastapi import FastAPI
@@ -55,22 +55,24 @@ class ServerDiscoveryUtil:
                 version = getattr(obj, "version", "unknown")
                 routes = ServerDiscoveryUtil._extract_route_info(obj)
                 route_count = len(routes)
-                
+
                 # Create a signature for deduplication
                 app_signature = (title, version, route_count)
-                
+
                 # Check if we've seen an identical app
                 if app_signature in seen_apps:
                     existing_app = seen_apps[app_signature]
                     # Compare route details to ensure they're truly identical
                     existing_routes = existing_app["routes"]
-                    
-                    if ServerDiscoveryUtil._routes_are_identical(routes, existing_routes):
+
+                    if ServerDiscoveryUtil._routes_are_identical(
+                        routes, existing_routes
+                    ):
                         logger.debug(
                             f"Skipping duplicate FastAPI app: '{title}' (same title, version, and routes)"
                         )
                         continue  # Skip this duplicate
-                
+
                 # This is a unique app, add it
                 app_id = f"app_{id(obj)}"
                 app_info = {
@@ -80,17 +82,19 @@ class ServerDiscoveryUtil:
                     "routes": routes,
                     "module": ServerDiscoveryUtil._get_app_module(obj),
                     "object_id": id(obj),
-                    "router_routes_count": len(obj.router.routes) if hasattr(obj, 'router') else 0,
+                    "router_routes_count": (
+                        len(obj.router.routes) if hasattr(obj, "router") else 0
+                    ),
                 }
-                
+
                 fastapi_apps[app_id] = app_info
                 seen_apps[app_signature] = app_info
-                
+
                 logger.debug(
                     f"Found FastAPI app: '{title}' (module: {app_info['module']}) with "
                     f"{len(routes)} routes"
                 )
-                
+
             except Exception as e:
                 logger.warning(f"Error analyzing FastAPI app: {e}")
                 continue
@@ -101,7 +105,7 @@ class ServerDiscoveryUtil:
     def discover_running_servers() -> List[Dict[str, Any]]:
         """
         Discover running uvicorn servers by scanning threads and checking port bindings.
-        
+
         Returns:
             List of server info dictionaries containing:
             - 'type': 'uvicorn' or 'unknown'
@@ -111,183 +115,208 @@ class ServerDiscoveryUtil:
             - 'app': FastAPI app if discoverable
         """
         running_servers = []
-        
+
         # Look for uvicorn server threads
         for thread in threading.enumerate():
-            if hasattr(thread, '_target'):
+            if hasattr(thread, "_target"):
                 # Check if thread target looks like a uvicorn server
-                target_name = getattr(thread._target, '__name__', '') if thread._target else ''
-                if 'server' in target_name.lower() or 'uvicorn' in target_name.lower():
+                target_name = (
+                    getattr(thread._target, "__name__", "") if thread._target else ""
+                )
+                if "server" in target_name.lower() or "uvicorn" in target_name.lower():
                     server_info = {
-                        'type': 'uvicorn',
-                        'thread': thread,
-                        'target_name': target_name,
-                        'daemon': thread.daemon,
-                        'alive': thread.is_alive(),
+                        "type": "uvicorn",
+                        "thread": thread,
+                        "target_name": target_name,
+                        "daemon": thread.daemon,
+                        "alive": thread.is_alive(),
                     }
-                    
+
                     # Try to extract server details from thread
-                    server_details = ServerDiscoveryUtil._extract_server_details_from_thread(thread)
+                    server_details = (
+                        ServerDiscoveryUtil._extract_server_details_from_thread(thread)
+                    )
                     server_info.update(server_details)
-                    
+
                     running_servers.append(server_info)
-                    logger.debug(f"Found running server thread: {target_name} (daemon={thread.daemon})")
-        
+                    logger.debug(
+                        f"Found running server thread: {target_name} (daemon={thread.daemon})"
+                    )
+
         # Also check for bound ports that might indicate running servers
         bound_ports = ServerDiscoveryUtil._discover_bound_ports()
         for port_info in bound_ports:
             # Only add if we haven't already found this port via thread discovery
-            existing_ports = [s.get('port') for s in running_servers if s.get('port')]
-            if port_info['port'] not in existing_ports:
-                port_info['type'] = 'unknown'
+            existing_ports = [s.get("port") for s in running_servers if s.get("port")]
+            if port_info["port"] not in existing_ports:
+                port_info["type"] = "unknown"
                 running_servers.append(port_info)
-                logger.debug(f"Found bound port: {port_info['host']}:{port_info['port']}")
-        
+                logger.debug(
+                    f"Found bound port: {port_info['host']}:{port_info['port']}"
+                )
+
         return running_servers
 
     @staticmethod
     def _extract_server_details_from_thread(thread) -> Dict[str, Any]:
         """Extract server details from a thread if possible."""
         details = {}
-        
+
         try:
             # Try to access thread local variables or target args
-            if hasattr(thread, '_args') and thread._args:
+            if hasattr(thread, "_args") and thread._args:
                 # Some uvicorn servers might have args with host/port
                 args = thread._args
                 if len(args) >= 2:
                     # Common pattern: (app, host, port) or similar
-                    if isinstance(args[0], str) and ':' in args[0]:
+                    if isinstance(args[0], str) and ":" in args[0]:
                         # Might be "host:port" format
                         try:
-                            host, port = args[0].split(':')
-                            details['host'] = host
-                            details['port'] = int(port)
+                            host, port = args[0].split(":")
+                            details["host"] = host
+                            details["port"] = int(port)
                         except (ValueError, IndexError):
                             pass
-            
+
             # Try to find FastAPI app in thread target or args
-            if hasattr(thread, '_target') and thread._target:
+            if hasattr(thread, "_target") and thread._target:
                 # Check if target has app attribute or if it's in closure
                 target = thread._target
-                if hasattr(target, '__closure__') and target.__closure__:
+                if hasattr(target, "__closure__") and target.__closure__:
                     for cell in target.__closure__:
                         try:
                             cell_contents = cell.cell_contents
                             from fastapi import FastAPI
+
                             if isinstance(cell_contents, FastAPI):
-                                details['app'] = cell_contents
-                                details['app_title'] = getattr(cell_contents, 'title', 'Unknown')
+                                details["app"] = cell_contents
+                                details["app_title"] = getattr(
+                                    cell_contents, "title", "Unknown"
+                                )
                                 break
                         except (ImportError, AttributeError):
                             continue
-                            
+
         except Exception as e:
             logger.debug(f"Could not extract server details from thread: {e}")
-        
+
         return details
 
     @staticmethod
     def _discover_bound_ports() -> List[Dict[str, Any]]:
         """Discover ports that are currently bound by this process."""
         bound_ports = []
-        
+
         try:
             # Common port ranges for web servers
             common_ports = [8000, 8080, 8090, 9090, 9091, 3000, 3001, 4000, 5000]
-            
+
             for port in common_ports:
-                for host in ['127.0.0.1', '0.0.0.0', 'localhost']:
+                for host in ["127.0.0.1", "0.0.0.0", "localhost"]:
                     try:
                         # Try to connect to see if port is bound
                         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                         sock.settimeout(0.1)  # Very short timeout
                         result = sock.connect_ex((host, port))
                         sock.close()
-                        
+
                         if result == 0:  # Connection successful = port is bound
-                            bound_ports.append({
-                                'host': host,
-                                'port': port,
-                                'status': 'bound'
-                            })
+                            bound_ports.append(
+                                {"host": host, "port": port, "status": "bound"}
+                            )
                             break  # Don't check other hosts for same port
                     except Exception:
                         continue
-                        
+
         except Exception as e:
             logger.debug(f"Error discovering bound ports: {e}")
-        
+
         return bound_ports
 
     @staticmethod
-    def find_server_on_port(target_port: int, target_host: str = None) -> Optional[Dict[str, Any]]:
+    def find_server_on_port(
+        target_port: int, target_host: str = None
+    ) -> Optional[Dict[str, Any]]:
         """
         Find if there's already a server running on the specified port.
-        
+
         Args:
             target_port: Port to check
             target_host: Host to check (optional)
-            
+
         Returns:
             Server info dict if found, None otherwise
         """
         running_servers = ServerDiscoveryUtil.discover_running_servers()
-        
+
         for server in running_servers:
-            server_port = server.get('port')
-            server_host = server.get('host')
-            
+            server_port = server.get("port")
+            server_host = server.get("host")
+
             # Check port match
             if server_port == target_port:
                 # If target_host is specified, check host match too
-                if target_host is None or server_host == target_host or server_host in ['0.0.0.0', '127.0.0.1']:
-                    logger.info(f"🔍 DISCOVERY: Found existing server on {server_host}:{server_port}")
+                if (
+                    target_host is None
+                    or server_host == target_host
+                    or server_host in ["0.0.0.0", "127.0.0.1"]
+                ):
+                    logger.info(
+                        f"🔍 DISCOVERY: Found existing server on {server_host}:{server_port}"
+                    )
                     return server
-        
+
         logger.debug(f"🔍 DISCOVERY: No existing server found on port {target_port}")
         return None
 
     @staticmethod
-    def _routes_are_identical(routes1: List[Dict[str, Any]], routes2: List[Dict[str, Any]]) -> bool:
+    def _routes_are_identical(
+        routes1: List[Dict[str, Any]], routes2: List[Dict[str, Any]]
+    ) -> bool:
         """Compare two route lists to see if they're identical."""
         if len(routes1) != len(routes2):
             return False
-            
+
         # Create comparable signatures for each route
         def route_signature(route):
             return (
-                tuple(sorted(route.get('methods', []))),  # Sort methods for consistent comparison
-                route.get('path', ''),
-                route.get('endpoint_name', '')
+                tuple(
+                    sorted(route.get("methods", []))
+                ),  # Sort methods for consistent comparison
+                route.get("path", ""),
+                route.get("endpoint_name", ""),
             )
-        
+
         # Sort routes by signature for consistent comparison
         sig1 = sorted([route_signature(r) for r in routes1])
         sig2 = sorted([route_signature(r) for r in routes2])
-        
+
         return sig1 == sig2
 
     @staticmethod
     def _extract_route_info(app) -> List[Dict[str, Any]]:
         """Extract route information from FastAPI app without modifying it."""
         routes = []
-        
+
         try:
             for route in app.router.routes:
-                if hasattr(route, 'endpoint') and hasattr(route, 'path'):
+                if hasattr(route, "endpoint") and hasattr(route, "path"):
                     route_info = {
                         "path": route.path,
-                        "methods": list(route.methods) if hasattr(route, 'methods') else [],
+                        "methods": (
+                            list(route.methods) if hasattr(route, "methods") else []
+                        ),
                         "endpoint": route.endpoint,
-                        "endpoint_name": getattr(route.endpoint, '__name__', 'unknown'),
-                        "has_mesh_route": hasattr(route.endpoint, '_mesh_route_metadata'),
+                        "endpoint_name": getattr(route.endpoint, "__name__", "unknown"),
+                        "has_mesh_route": hasattr(
+                            route.endpoint, "_mesh_route_metadata"
+                        ),
                     }
                     routes.append(route_info)
-                    
+
         except Exception as e:
             logger.warning(f"Error extracting route info: {e}")
-            
+
         return routes
 
     @staticmethod
@@ -297,16 +326,16 @@ class ServerDiscoveryUtil:
             # Try to get module from the app's stack frame when it was created
             # This is best-effort - may not always work
             import inspect
-            
+
             frame = inspect.currentframe()
             while frame:
                 frame_globals = frame.f_globals
                 for name, obj in frame_globals.items():
                     if obj is app:
-                        return frame_globals.get('__name__', 'unknown')
+                        return frame_globals.get("__name__", "unknown")
                 frame = frame.f_back
-                
+
         except Exception:
             pass
-            
-        return None
+
+        return None

_mcp_mesh/shared/simple_shutdown.py

@@ -2,9 +2,9 @@
 Simple shutdown coordination for MCP Mesh agents.
 
 Provides clean shutdown via FastAPI lifespan events and basic signal handling.
+The Rust core handles actual deregistration from the registry.
 """
 
-import asyncio
 import logging
 import signal
 from contextlib import asynccontextmanager
@@ -14,7 +14,12 @@ logger = logging.getLogger(__name__)
 
 
 class SimpleShutdownCoordinator:
-    """Lightweight shutdown coordination using FastAPI lifespan."""
+    """Lightweight shutdown coordination using FastAPI lifespan.
+
+    The Rust core handles registry deregistration automatically when
+    handle.shutdown() is called. This coordinator just manages the
+    shutdown signal flow between Python and Rust.
+    """
 
     def __init__(self):
         self._shutdown_requested = False
@@ -23,7 +28,7 @@ class SimpleShutdownCoordinator:
         self._shutdown_complete = False  # Flag to prevent race conditions
 
     def set_shutdown_context(self, registry_url: str, agent_id: str) -> None:
-        """Set context for shutdown cleanup."""
+        """Set context for shutdown (used for logging)."""
         self._registry_url = registry_url
         self._agent_id = agent_id
         logger.debug(
@@ -54,61 +59,20 @@ class SimpleShutdownCoordinator:
         self._shutdown_complete = True
         logger.debug("🏁 Shutdown marked as complete")
 
-    async def perform_registry_cleanup(self) -> None:
-        """Perform registry cleanup by calling DELETE /agents/{agent_id}."""
-        # Try to get the actual agent_id from DecoratorRegistry if available
-        actual_agent_id = self._agent_id
-        try:
-            from _mcp_mesh.engine.decorator_registry import DecoratorRegistry
-
-            agent_config = DecoratorRegistry.get_resolved_agent_config()
-            if agent_config and "agent_id" in agent_config:
-                resolved_agent_id = agent_config["agent_id"]
-                if resolved_agent_id and resolved_agent_id != "unknown":
-                    actual_agent_id = resolved_agent_id
-                    logger.debug(
-                        f"🔧 Using resolved agent_id from DecoratorRegistry: {actual_agent_id}"
-                    )
-        except Exception as e:
-            logger.debug(f"Could not get agent_id from DecoratorRegistry: {e}")
-
-        if (
-            not self._registry_url
-            or not actual_agent_id
-            or actual_agent_id == "unknown"
-        ):
-            logger.warning(
-                f"⚠️ Missing registry URL or agent ID for cleanup: registry_url={self._registry_url}, agent_id={actual_agent_id}"
-            )
-            return
-
-        try:
-            from _mcp_mesh.generated.mcp_mesh_registry_client.api_client import (
-                ApiClient,
-            )
-            from _mcp_mesh.generated.mcp_mesh_registry_client.configuration import (
-                Configuration,
-            )
-            from _mcp_mesh.shared.registry_client_wrapper import RegistryClientWrapper
-
-            config = Configuration(host=self._registry_url)
-            api_client = ApiClient(configuration=config)
-            registry_wrapper = RegistryClientWrapper(api_client)
-
-            success = await registry_wrapper.unregister_agent(actual_agent_id)
-            if success:
-                logger.info(f"✅ Agent '{actual_agent_id}' unregistered from registry")
-                self.mark_shutdown_complete()
-            else:
-                logger.warning(f"⚠️ Failed to unregister agent '{actual_agent_id}'")
-                self.mark_shutdown_complete()  # Mark complete even on failure to prevent loops
-
-        except Exception as e:
-            logger.error(f"❌ Registry cleanup error: {e}")
-            self.mark_shutdown_complete()  # Mark complete even on error to prevent loops
+    def request_shutdown(self) -> None:
+        """Request shutdown (called when lifespan exits)."""
+        self._shutdown_requested = True
+        agent_id = self._agent_id or "<unknown>"
+        logger.info(f"🔄 Shutdown requested for agent '{agent_id}'")
 
     def create_shutdown_lifespan(self, original_lifespan=None):
-        """Create lifespan function that includes registry cleanup."""
+        """Create lifespan function that signals shutdown on exit.
+
+        The Rust core will handle actual deregistration when it receives
+        the shutdown signal via handle.shutdown().
+        """
+        # Capture agent_id at creation time with fallback for None
+        agent_id = self._agent_id or "<unknown>"
 
         @asynccontextmanager
         async def shutdown_lifespan(app):
@@ -120,10 +84,14 @@ class SimpleShutdownCoordinator:
             else:
                 yield
 
-            # Shutdown phase
-            logger.info("🔄 FastAPI shutdown initiated, performing registry cleanup...")
-            await self.perform_registry_cleanup()
-            logger.info("🏁 Registry cleanup completed")
+            # Shutdown phase - just signal, Rust handles deregistration
+            logger.info(
+                f"🔄 FastAPI shutdown initiated for agent '{agent_id}', "
+                "Rust core will handle deregistration"
+            )
+            self.request_shutdown()
+            self.mark_shutdown_complete()
+            logger.info("🏁 Shutdown signaled")
 
         return shutdown_lifespan
 
@@ -164,12 +132,15 @@ def start_blocking_loop_with_shutdown_support(thread) -> None:
     """
     Keep main thread alive while uvicorn in the thread handles requests.
 
-    Install signal handlers in main thread for proper registry cleanup since
+    Install signal handlers in main thread for proper shutdown signaling since
     signals to threads can be unreliable for FastAPI lifespan shutdown.
+
+    Note: The Rust core handles registry deregistration automatically when
+    handle.shutdown() is called from the heartbeat task.
     """
-    logger.info("🔒 MAIN THREAD: Installing signal handlers for registry cleanup")
+    logger.info("🔒 MAIN THREAD: Installing signal handlers")
 
-    # Install signal handlers for proper registry cleanup
+    # Install signal handlers
     _simple_shutdown_coordinator.install_signal_handlers()
 
     logger.info(
@@ -184,34 +155,21 @@ def start_blocking_loop_with_shutdown_support(thread) -> None:
             # Check if shutdown was requested via signal
             if _simple_shutdown_coordinator.is_shutdown_requested():
                 logger.info(
-                    "🔄 MAIN THREAD: Shutdown requested, performing registry cleanup..."
+                    "🔄 MAIN THREAD: Shutdown requested, signaling heartbeat to stop..."
+                )
+                # Mark shutdown complete so heartbeat task will call handle.shutdown()
+                # which triggers Rust core to deregister from registry
+                _simple_shutdown_coordinator.mark_shutdown_complete()
+                logger.info(
+                    "🏁 MAIN THREAD: Shutdown signaled, Rust core will handle deregistration"
                 )
-
-                # Perform registry cleanup in main thread
-                import asyncio
-
-                try:
-                    # Run cleanup in main thread
-                    asyncio.run(_simple_shutdown_coordinator.perform_registry_cleanup())
-                except Exception as e:
-                    logger.error(f"❌ Registry cleanup error: {e}")
-
-                logger.info("🏁 MAIN THREAD: Registry cleanup completed, exiting")
                 break
 
     except KeyboardInterrupt:
+        logger.info("🔄 MAIN THREAD: KeyboardInterrupt received, signaling shutdown...")
+        _simple_shutdown_coordinator.mark_shutdown_complete()
         logger.info(
-            "🔄 MAIN THREAD: KeyboardInterrupt received, performing registry cleanup..."
+            "🏁 MAIN THREAD: Shutdown signaled, Rust core will handle deregistration"
         )
 
-        # Perform registry cleanup on Ctrl+C
-        import asyncio
-
-        try:
-            asyncio.run(_simple_shutdown_coordinator.perform_registry_cleanup())
-        except Exception as e:
-            logger.error(f"❌ Registry cleanup error: {e}")
-
-        logger.info("🏁 MAIN THREAD: Registry cleanup completed")
-
     logger.info("🏁 MAIN THREAD: Uvicorn thread completed")

_mcp_mesh/tracing/execution_tracer.py

@@ -10,12 +10,8 @@ from collections.abc import Callable
 from typing import Any, Optional
 
 # Import shared utilities at module level to avoid circular imports during execution
-from .utils import (
-    generate_span_id,
-    get_agent_metadata_with_fallback,
-    is_tracing_enabled,
-    publish_trace_with_fallback,
-)
+from .utils import (generate_span_id, get_agent_metadata_with_fallback,
+                    is_tracing_enabled, publish_trace_with_fallback)
 
 logger = logging.getLogger(__name__)