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

_mcp_mesh/engine/signature_analyzer.py

@@ -5,21 +5,67 @@ Function signature analysis for MCP Mesh dependency injection.
 import inspect
 from typing import Any, get_type_hints
 
-from mesh.types import McpMeshAgent, MeshLlmAgent
+from mesh.types import McpMeshTool, MeshLlmAgent
+
+# Also support deprecated McpMeshAgent for backwards compatibility
+try:
+    from mesh.types import McpMeshAgent
+except ImportError:
+    McpMeshAgent = McpMeshTool  # type: ignore
+
+
+def _is_mesh_tool_type(param_type: Any) -> bool:
+    """Check if a type is McpMeshTool or deprecated McpMeshAgent."""
+    # Direct McpMeshTool type
+    if (
+        param_type == McpMeshTool
+        or (hasattr(param_type, "__name__") and param_type.__name__ == "McpMeshTool")
+        or (
+            hasattr(param_type, "__origin__")
+            and param_type.__origin__ == type(McpMeshTool)
+        )
+    ):
+        return True
+
+    # Support deprecated McpMeshAgent
+    if (
+        param_type == McpMeshAgent
+        or (hasattr(param_type, "__name__") and param_type.__name__ == "McpMeshAgent")
+        or (
+            hasattr(param_type, "__origin__")
+            and param_type.__origin__ == type(McpMeshAgent)
+        )
+    ):
+        return True
+
+    # Union type (e.g., McpMeshTool | None)
+    if hasattr(param_type, "__args__"):
+        for arg in param_type.__args__:
+            if arg == McpMeshTool or (
+                hasattr(arg, "__name__") and arg.__name__ == "McpMeshTool"
+            ):
+                return True
+            # Support deprecated McpMeshAgent in unions
+            if arg == McpMeshAgent or (
+                hasattr(arg, "__name__") and arg.__name__ == "McpMeshAgent"
+            ):
+                return True
+
+    return False
 
 
 def get_mesh_agent_positions(func: Any) -> list[int]:
     """
-    Get positions of McpMeshAgent parameters in function signature.
+    Get positions of McpMeshTool parameters in function signature.
 
     Args:
         func: Function to analyze
 
     Returns:
-        List of parameter positions (0-indexed) that are McpMeshAgent types
+        List of parameter positions (0-indexed) that are McpMeshTool types
 
     Example:
-        def greet(name: str, date_svc: McpMeshAgent, file_svc: McpMeshAgent):
+        def greet(name: str, date_svc: McpMeshTool, file_svc: McpMeshTool):
             pass
 
         get_mesh_agent_positions(greet) → [1, 2]
@@ -32,40 +78,12 @@ def get_mesh_agent_positions(func: Any) -> list[int]:
         sig = inspect.signature(func)
         param_names = list(sig.parameters.keys())
 
-        # Find positions of McpMeshAgent parameters
+        # Find positions of McpMeshTool parameters
         mesh_positions = []
         for i, param_name in enumerate(param_names):
             if param_name in type_hints:
                 param_type = type_hints[param_name]
-
-                # Check if it's McpMeshAgent type (handle different import paths and Union types)
-                is_agent = False
-
-                # Direct McpMeshAgent type
-                if (
-                    param_type == McpMeshAgent
-                    or (
-                        hasattr(param_type, "__name__")
-                        and param_type.__name__ == "McpMeshAgent"
-                    )
-                    or (
-                        hasattr(param_type, "__origin__")
-                        and param_type.__origin__ == type(McpMeshAgent)
-                    )
-                ):
-                    is_agent = True
-
-                # Union type (e.g., McpMeshAgent | None)
-                elif hasattr(param_type, "__args__"):
-                    # Check if any arg in the union is McpMeshAgent
-                    for arg in param_type.__args__:
-                        if arg == McpMeshAgent or (
-                            hasattr(arg, "__name__") and arg.__name__ == "McpMeshAgent"
-                        ):
-                            is_agent = True
-                            break
-
-                if is_agent:
+                if _is_mesh_tool_type(param_type):
                     mesh_positions.append(i)
 
         return mesh_positions
@@ -81,13 +99,13 @@ def get_mesh_agent_positions(func: Any) -> list[int]:
 
 def get_mesh_agent_parameter_names(func: Any) -> list[str]:
     """
-    Get names of McpMeshAgent parameters in function signature.
+    Get names of McpMeshTool parameters in function signature.
 
     Args:
         func: Function to analyze
 
     Returns:
-        List of parameter names that are McpMeshAgent types
+        List of parameter names that are McpMeshTool types
     """
     try:
         type_hints = get_type_hints(func)
@@ -97,35 +115,7 @@ def get_mesh_agent_parameter_names(func: Any) -> list[str]:
         for param_name, param in sig.parameters.items():
             if param_name in type_hints:
                 param_type = type_hints[param_name]
-
-                # Check if it's McpMeshAgent type (handle different import paths and Union types)
-                is_mesh_agent = False
-
-                # Direct McpMeshAgent type
-                if (
-                    param_type == McpMeshAgent
-                    or (
-                        hasattr(param_type, "__name__")
-                        and param_type.__name__ == "McpMeshAgent"
-                    )
-                    or (
-                        hasattr(param_type, "__origin__")
-                        and param_type.__origin__ == type(McpMeshAgent)
-                    )
-                ):
-                    is_mesh_agent = True
-
-                # Union type (e.g., McpMeshAgent | None)
-                elif hasattr(param_type, "__args__"):
-                    # Check if any arg in the union is McpMeshAgent
-                    for arg in param_type.__args__:
-                        if arg == McpMeshAgent or (
-                            hasattr(arg, "__name__") and arg.__name__ == "McpMeshAgent"
-                        ):
-                            is_mesh_agent = True
-                            break
-
-                if is_mesh_agent:
+                if _is_mesh_tool_type(param_type):
                     mesh_param_names.append(param_name)
 
         return mesh_param_names
@@ -136,7 +126,7 @@ def get_mesh_agent_parameter_names(func: Any) -> list[str]:
 
 def validate_mesh_dependencies(func: Any, dependencies: list[dict]) -> tuple[bool, str]:
     """
-    Validate that the number of dependencies matches McpMeshAgent parameters.
+    Validate that the number of dependencies matches McpMeshTool parameters.
 
     Args:
         func: Function to validate
@@ -149,9 +139,9 @@ def validate_mesh_dependencies(func: Any, dependencies: list[dict]) -> tuple[boo
 
     if len(dependencies) != len(mesh_positions):
         return False, (
-            f"Function {func.__name__} has {len(mesh_positions)} McpMeshAgent parameters "
+            f"Function {func.__name__} has {len(mesh_positions)} McpMeshTool parameters "
             f"but {len(dependencies)} dependencies declared. "
-            f"Each McpMeshAgent parameter needs a corresponding dependency."
+            f"Each McpMeshTool parameter needs a corresponding dependency."
         )
 
     return True, ""

_mcp_mesh/engine/unified_mcp_proxy.py

@@ -16,6 +16,8 @@ from ..shared.logging_config import (
     get_trace_prefix,
 )
 from ..shared.sse_parser import SSEParser
+from ..tracing.context import TraceContext
+from ..tracing.utils import generate_span_id
 
 logger = logging.getLogger(__name__)
 
@@ -23,7 +25,7 @@ logger = logging.getLogger(__name__)
 class UnifiedMCPProxy:
     """Unified MCP proxy using FastMCP's built-in client.
 
-    This provides the implementation for McpMeshAgent type parameters,
+    This provides the implementation for McpMeshTool type parameters,
     offering all MCP protocol features using FastMCP's superior client.
 
     Features:
@@ -178,37 +180,6 @@ class UnifiedMCPProxy:
                 f"FastMCP client failed: {e}"
             )  # Convert to ImportError to trigger fallback
 
-    def _get_trace_headers(self) -> dict[str, str]:
-        """Extract trace headers from current context for distributed tracing.
-
-        Returns:
-            Dict of trace headers or empty dict if no trace context available
-        """
-        try:
-            from ..tracing.context import TraceContext
-
-            current_trace = TraceContext.get_current()
-            if current_trace:
-                headers = {
-                    "X-Trace-ID": current_trace.trace_id,
-                    "X-Parent-Span": current_trace.span_id,  # Current span becomes parent for downstream
-                }
-                self.logger.info(
-                    f"🔗 TRACE_PROPAGATION: Injecting headers trace_id={current_trace.trace_id[:8]}... "
-                    f"parent_span={current_trace.span_id[:8]}..."
-                )
-                return headers
-            else:
-                self.logger.warning("🔗 TRACE_PROPAGATION: No trace context available")
-                return {}
-
-        except Exception as e:
-            # Never fail MCP calls due to tracing issues
-            self.logger.warning(
-                f"🔗 TRACE_PROPAGATION: Exception getting trace context: {e}"
-            )
-            return {}
-
     def _configure_from_kwargs(self):
         """Auto-configure proxy settings from kwargs."""
         # Basic configuration
@@ -433,6 +404,19 @@ class UnifiedMCPProxy:
         # Get trace prefix if available
         tp = get_trace_prefix()
 
+        # Inject trace context into arguments for downstream agents
+        # This is the fallback mechanism for agents that can't access HTTP headers (e.g., TypeScript)
+        args_with_trace = dict(arguments) if arguments else {}
+        current_trace = TraceContext.get_current()
+        if current_trace:
+            # Use current function's span_id as parent for downstream call
+            # Don't generate a new span - that creates unpublished "ghost" spans that break the tree
+            args_with_trace["_trace_id"] = current_trace.trace_id
+            args_with_trace["_parent_span"] = current_trace.span_id
+            self.logger.debug(
+                f"{tp}🔗 Injecting trace context: trace_id={current_trace.trace_id[:8]}..., parent_span={current_trace.span_id[:8]}..."
+            )
+
         # Log cross-agent call - summary line
         arg_keys = list(arguments.keys()) if arguments else []
         self.logger.debug(
@@ -451,7 +435,7 @@ class UnifiedMCPProxy:
             async with client_instance as client:
 
                 # Use FastMCP's call_tool which returns CallToolResult object
-                result = await client.call_tool(name, arguments or {})
+                result = await client.call_tool(name, args_with_trace)
 
                 # Calculate performance metrics
                 end_time = time.time()
@@ -479,12 +463,12 @@ class UnifiedMCPProxy:
             self.logger.warning(
                 f"FastMCP Client not available: {e}, falling back to HTTP"
             )
-            return await self._fallback_http_call(name, arguments)
+            return await self._fallback_http_call(name, args_with_trace)
         except Exception as e:
             self.logger.warning(f"FastMCP Client failed: {e}, falling back to HTTP")
             # Try HTTP fallback
             try:
-                result = await self._fallback_http_call(name, arguments)
+                result = await self._fallback_http_call(name, args_with_trace)
                 return result
             except Exception as fallback_error:
                 raise RuntimeError(

_mcp_mesh/pipeline/__init__.py

@@ -2,31 +2,22 @@
 MCP Mesh Pipeline Architecture
 
 This module provides a clean, explicit pipeline-based architecture for processing
-decorators and managing the mesh agent lifecycle. This replaces the scattered
-async processing with a clear, sequential flow that can be easily tested and debugged.
+decorators and managing the mesh agent lifecycle. The Rust core handles registry
+communication including heartbeats, dependency resolution, and deregistration.
 
 Key Components:
 - MeshPipeline: Main orchestrator that executes steps in sequence
 - PipelineStep: Interface for individual processing steps
 - PipelineResult: Result container with status and context
-- Built-in steps for common operations (collection, config, heartbeat, etc.)
+- Built-in steps for common operations (collection, config, etc.)
 """
 
-from .mcp_heartbeat import (
-                        DependencyResolutionStep,
-                        HeartbeatSendStep,
-                        RegistryConnectionStep,
-)
+from .mcp_heartbeat import rust_heartbeat_task
+from .mcp_startup import (ConfigurationStep, DecoratorCollectionStep,
+                          FastAPIServerSetupStep, FastMCPServerDiscoveryStep,
+                          HeartbeatLoopStep, HeartbeatPreparationStep,
+                          StartupPipeline)
 from .shared import MeshPipeline, PipelineResult, PipelineStatus, PipelineStep
-from .mcp_startup import (
-                        ConfigurationStep,
-                        DecoratorCollectionStep,
-                        FastAPIServerSetupStep,
-                        FastMCPServerDiscoveryStep,
-                        HeartbeatLoopStep,
-                        HeartbeatPreparationStep,
-                        StartupPipeline,
-)
 
 __all__ = [
     "MeshPipeline",
@@ -39,8 +30,6 @@ __all__ = [
     "FastMCPServerDiscoveryStep",
     "HeartbeatLoopStep",
     "HeartbeatPreparationStep",
-    "RegistryConnectionStep",
-    "HeartbeatSendStep",
-    "DependencyResolutionStep",
     "StartupPipeline",
+    "rust_heartbeat_task",
 ]

_mcp_mesh/pipeline/api_heartbeat/__init__.py

@@ -3,14 +3,19 @@ API heartbeat pipeline for FastAPI integration.
 
 Provides periodic service registration and health monitoring
 for FastAPI applications using @mesh.route decorators.
+
+Uses Rust core for registry communication, dependency resolution,
+and deregistration.
 """
 
-from .api_heartbeat_pipeline import APIHeartbeatPipeline
-from .api_heartbeat_orchestrator import APIHeartbeatOrchestrator
-from .api_dependency_resolution import APIDependencyResolutionStep
+from .api_lifespan_integration import (api_heartbeat_lifespan_task,
+                                       create_api_lifespan_handler,
+                                       integrate_api_heartbeat_with_fastapi)
+from .rust_api_heartbeat import rust_api_heartbeat_task
 
 __all__ = [
-    "APIHeartbeatPipeline", 
-    "APIHeartbeatOrchestrator",
-    "APIDependencyResolutionStep",
-]
+    "api_heartbeat_lifespan_task",
+    "create_api_lifespan_handler",
+    "integrate_api_heartbeat_with_fastapi",
+    "rust_api_heartbeat_task",
+]

_mcp_mesh/pipeline/api_heartbeat/api_lifespan_integration.py

@@ -1,8 +1,10 @@
 """
 FastAPI lifespan integration for API heartbeat pipeline.
 
-Handles the execution of API heartbeat pipeline as a background task
+Handles the execution of API heartbeat as a background task
 during FastAPI application lifespan for @mesh.route decorator services.
+
+Uses the Rust core for registry communication.
 """
 
 import asyncio
@@ -14,65 +16,35 @@ logger = logging.getLogger(__name__)
 
 async def api_heartbeat_lifespan_task(heartbeat_config: dict[str, Any]) -> None:
     """
-    API heartbeat task that runs in FastAPI lifespan using pipeline architecture.
+    API heartbeat task that runs in FastAPI lifespan.
+
+    Uses Rust-backed heartbeat for registry communication.
 
     Args:
-        heartbeat_config: Configuration containing service_id, interval, 
+        heartbeat_config: Configuration containing service_id, interval,
                          and context for API heartbeat execution
     """
-    service_id = heartbeat_config["service_id"]
-    interval = heartbeat_config["interval"]  # Already validated by get_config_value in setup
-    context = heartbeat_config["context"]
+    service_id = heartbeat_config.get("service_id", "unknown-api-service")
     standalone_mode = heartbeat_config.get("standalone_mode", False)
 
     # Check if running in standalone mode
     if standalone_mode:
         logger.info(
-            f"💓 Starting API heartbeat pipeline in standalone mode for service '{service_id}' "
+            f"💓 API heartbeat in standalone mode for service '{service_id}' "
             f"(no registry communication)"
         )
-        return  # For now, skip heartbeat in standalone mode
-
-    # Create API heartbeat orchestrator for pipeline execution
-    from .api_heartbeat_orchestrator import APIHeartbeatOrchestrator
+        return
 
-    api_heartbeat_orchestrator = APIHeartbeatOrchestrator()
+    # Use Rust-backed heartbeat
+    from .rust_api_heartbeat import rust_api_heartbeat_task
 
-    logger.info(f"💓 Starting API heartbeat pipeline task for service '{service_id}'")
-
-    try:
-        while True:
-            try:
-                # Execute API heartbeat pipeline
-                success = await api_heartbeat_orchestrator.execute_api_heartbeat(
-                    service_id, context
-                )
-
-                if not success:
-                    # Log failure but continue to next cycle (pipeline handles detailed logging)
-                    logger.debug(
-                        f"💔 API heartbeat pipeline failed for service '{service_id}' - "
-                        f"continuing to next cycle"
-                    )
-
-            except Exception as e:
-                # Log pipeline execution error but continue to next cycle for resilience
-                logger.error(
-                    f"❌ API heartbeat pipeline execution error for service '{service_id}': {e}"
-                )
-                # Continue to next cycle - heartbeat should be resilient
-
-            # Wait for next heartbeat interval
-            await asyncio.sleep(interval)
-
-    except asyncio.CancelledError:
-        logger.info(f"🛑 API heartbeat pipeline task cancelled for service '{service_id}'")
-        raise
+    logger.info(f"💓 Using Rust-backed heartbeat for API service '{service_id}'")
+    await rust_api_heartbeat_task(heartbeat_config)
 
 
 def create_api_lifespan_handler(heartbeat_config: dict[str, Any]) -> Any:
     """
-    Create a FastAPI lifespan context manager that runs API heartbeat pipeline.
+    Create a FastAPI lifespan context manager that runs API heartbeat.
 
     Args:
         heartbeat_config: Configuration for API heartbeat execution
@@ -100,11 +72,13 @@ def create_api_lifespan_handler(heartbeat_config: dict[str, Any]) -> Any:
             # Cleanup: cancel heartbeat task
             logger.info(f"🛑 Shutting down FastAPI lifespan for service '{service_id}'")
             heartbeat_task.cancel()
-            
+
             try:
                 await heartbeat_task
             except asyncio.CancelledError:
-                logger.info(f"✅ API heartbeat task cancelled for service '{service_id}'")
+                logger.info(
+                    f"✅ API heartbeat task cancelled for service '{service_id}'"
+                )
 
     return api_lifespan
 
@@ -113,18 +87,18 @@ def integrate_api_heartbeat_with_fastapi(
     fastapi_app: Any, heartbeat_config: dict[str, Any]
 ) -> None:
     """
-    Integrate API heartbeat pipeline with FastAPI lifespan events.
+    Integrate API heartbeat with FastAPI lifespan events.
 
     Args:
         fastapi_app: FastAPI application instance
         heartbeat_config: Configuration for heartbeat execution
     """
     service_id = heartbeat_config.get("service_id", "unknown")
-    
+
     try:
         # Check if FastAPI app already has a lifespan handler
         existing_lifespan = getattr(fastapi_app, "router.lifespan_context", None)
-        
+
         if existing_lifespan is not None:
             logger.warning(
                 f"⚠️ FastAPI app already has lifespan handler - "
@@ -144,4 +118,4 @@ def integrate_api_heartbeat_with_fastapi(
             f"❌ Failed to integrate API heartbeat with FastAPI lifespan "
             f"for service '{service_id}': {e}"
         )
-        raise
+        raise