mcp-mesh 0.5.1__py3-none-any.whl → 0.5.3__py3-none-any.whl

_mcp_mesh/__init__.py

@@ -31,7 +31,7 @@ from .engine.decorator_registry import (
     get_decorator_stats,
 )
 
-__version__ = "0.5.1"
+__version__ = "0.5.3"
 
 # Store reference to runtime processor if initialized
 _runtime_processor = None

_mcp_mesh/engine/dependency_injector.py

@@ -2,6 +2,8 @@
 Dynamic dependency injection system for MCP Mesh.
 
 Handles both initial injection and runtime updates when topology changes.
+Focused purely on dependency injection - telemetry/tracing is handled at 
+the HTTP middleware layer for unified approach across MCP agents and FastAPI apps.
 """
 
 import asyncio
@@ -292,41 +294,23 @@ class DependencyInjector:
             if inspect.iscoroutinefunction(func):
                 @functools.wraps(func)
                 async def minimal_wrapper(*args, **kwargs):
-                    # Execute with telemetry tracing even for async functions without dependencies
-                    try:
-                        from ..tracing.execution_tracer import ExecutionTracer
-                        
-                        # Use ExecutionTracer for minimal async wrapper (no dependencies)
-                        return await ExecutionTracer.trace_function_execution_async(
-                            func, args, kwargs, [], [], 0, wrapper_logger
-                        )
-                    except ImportError:
-                        # Fallback if tracing is unavailable - never fail user function
-                        wrapper_logger.debug("🔇 Tracing unavailable, executing without telemetry")
-                        return await func(*args, **kwargs)
-                    except Exception as e:
-                        # Never fail user function due to tracing errors
-                        wrapper_logger.warning(f"⚠️ Telemetry failed, executing without: {e}")
-                        return await func(*args, **kwargs)
+                    # Use ExecutionTracer for functions without dependencies (v0.4.0 style)
+                    from ..tracing.execution_tracer import ExecutionTracer
+                    wrapper_logger.debug(f"🔧 DI: Executing async function {func.__name__} (no dependencies)")
+                    
+                    # For async functions without dependencies, use the async tracer
+                    return await ExecutionTracer.trace_function_execution_async(
+                        func, args, kwargs, [], [], 0, wrapper_logger
+                    )
             else:
                 @functools.wraps(func)
                 def minimal_wrapper(*args, **kwargs):
-                    # Execute with telemetry tracing even for functions without dependencies
-                    try:
-                        from ..tracing.execution_tracer import ExecutionTracer
-                        
-                        # Use ExecutionTracer for minimal wrapper (no dependencies)
-                        return ExecutionTracer.trace_original_function(
-                            func, args, kwargs, wrapper_logger
-                        )
-                    except ImportError:
-                        # Fallback if tracing is unavailable - never fail user function
-                        wrapper_logger.debug("🔇 Tracing unavailable, executing without telemetry")
-                        return func(*args, **kwargs)
-                    except Exception as e:
-                        # Never fail user function due to tracing errors
-                        wrapper_logger.warning(f"⚠️ Telemetry failed, executing without: {e}")
-                        return func(*args, **kwargs)
+                    # Use ExecutionTracer for functions without dependencies (v0.4.0 style)
+                    from ..tracing.execution_tracer import ExecutionTracer
+                    wrapper_logger.debug(f"🔧 DI: Executing sync function {func.__name__} (no dependencies)")
+                    
+                    # Use original function tracer for functions without dependencies
+                    return ExecutionTracer.trace_original_function(func, args, kwargs, wrapper_logger)
 
             # Add minimal metadata for compatibility
             minimal_wrapper._mesh_injected_deps = {}
@@ -426,54 +410,29 @@ class DependencyInjector:
                     f"🔧 DEPENDENCY_WRAPPER: final_kwargs={final_kwargs}"
                 )
 
-                # ===== EXECUTE WITH TELEMETRY AND DEPENDENCY INJECTION =====
-                # Call the original function with telemetry tracing and injected dependencies
+                # ===== EXECUTE WITH DEPENDENCY INJECTION AND TRACING =====
+                # Use ExecutionTracer for comprehensive execution logging (v0.4.0 style)
+                from ..tracing.execution_tracer import ExecutionTracer
+                
                 original_func = func._mesh_original_func
 
-                # Execute with telemetry tracing
-                try:
-                    from ..tracing.execution_tracer import ExecutionTracer
-                    
-                    # Use ExecutionTracer for comprehensive telemetry
-                    if inspect.iscoroutinefunction(original_func):
-                        # For async functions, await the result directly
-                        result = await ExecutionTracer.trace_function_execution_async(
-                            original_func,
-                            args,
-                            final_kwargs,
-                            dependencies,
-                            mesh_positions,
-                            injected_count,
-                            wrapper_logger,
-                        )
-                    else:
-                        # For sync functions, call directly
-                        result = ExecutionTracer.trace_function_execution(
-                            original_func,
-                            args,
-                            final_kwargs,
-                            dependencies,
-                            mesh_positions,
-                            injected_count,
-                            wrapper_logger,
-                        )
-                except ImportError:
-                    # Fallback if tracing is unavailable - never fail user function
-                    wrapper_logger.debug("🔇 Tracing unavailable, executing without telemetry")
-                    if inspect.iscoroutinefunction(original_func):
-                        result = await original_func(*args, **final_kwargs)
-                    else:
-                        result = original_func(*args, **final_kwargs)
-                except Exception as e:
-                    # Never fail user function due to tracing errors
-                    wrapper_logger.warning(f"⚠️ Telemetry failed, executing without: {e}")
-                    if inspect.iscoroutinefunction(original_func):
-                        result = await original_func(*args, **final_kwargs)
-                    else:
-                        result = original_func(*args, **final_kwargs)
+                wrapper_logger.debug(
+                    f"🔧 DI: Executing async function {original_func.__name__} with {injected_count} injected dependencies"
+                )
+
+                # Use ExecutionTracer's async method for clean tracing
+                result = await ExecutionTracer.trace_function_execution_async(
+                    original_func,
+                    args,
+                    final_kwargs,
+                    dependencies,
+                    mesh_positions,
+                    injected_count,
+                    wrapper_logger,
+                )
 
                 wrapper_logger.debug(
-                    f"🔧 DEPENDENCY_WRAPPER: Original returned: {type(result)}"
+                    f"🔧 DI: Function {original_func.__name__} returned: {type(result)}"
                 )
                 return result
 
@@ -524,29 +483,24 @@ class DependencyInjector:
                             final_kwargs[param_name] = dependency
                             injected_count += 1
 
-                # ===== EXECUTE WITH TELEMETRY AND DEPENDENCY INJECTION =====
-                # Call the original function with telemetry tracing and injected dependencies
-                try:
-                    from ..tracing.execution_tracer import ExecutionTracer
-                    
-                    # Use ExecutionTracer for comprehensive telemetry
-                    return ExecutionTracer.trace_function_execution(
-                        func._mesh_original_func,
-                        args,
-                        final_kwargs,
-                        dependencies,
-                        mesh_positions,
-                        injected_count,
-                        wrapper_logger,
-                    )
-                except ImportError:
-                    # Fallback if tracing is unavailable - never fail user function
-                    wrapper_logger.debug("🔇 Tracing unavailable, executing without telemetry")
-                    return func._mesh_original_func(*args, **final_kwargs)
-                except Exception as e:
-                    # Never fail user function due to tracing errors
-                    wrapper_logger.warning(f"⚠️ Telemetry failed, executing without: {e}")
-                    return func._mesh_original_func(*args, **final_kwargs)
+                # ===== EXECUTE WITH DEPENDENCY INJECTION AND TRACING =====
+                # Use ExecutionTracer for comprehensive execution logging (v0.4.0 style)
+                from ..tracing.execution_tracer import ExecutionTracer
+                
+                wrapper_logger.debug(
+                    f"🔧 DI: Executing sync function {func._mesh_original_func.__name__} with {injected_count} injected dependencies"
+                )
+
+                # Use ExecutionTracer for clean execution tracing
+                return ExecutionTracer.trace_function_execution(
+                    func._mesh_original_func,
+                    args,
+                    final_kwargs,
+                    dependencies,
+                    mesh_positions,
+                    injected_count,
+                    wrapper_logger,
+                )
 
         # Store dependency state on wrapper
         dependency_wrapper._mesh_injected_deps = {}

_mcp_mesh/engine/http_wrapper.py

@@ -395,11 +395,10 @@ class HttpMcpWrapper:
         from starlette.responses import Response
 
         class MCPSessionRoutingMiddleware(BaseHTTPMiddleware):
-            """Session routing middleware for MCP requests.
+            """Clean session routing middleware for MCP requests (v0.4.0 style).
             
-            Handles session affinity by routing requests to appropriate pods
-            based on session ID. Telemetry/tracing is now handled in the
-            DI function wrapper for unified coverage of both MCP and API calls.
+            Handles session affinity and basic trace context setup only.
+            Function execution tracing is handled by ExecutionTracer in DependencyInjector.
             """
             def __init__(self, app, http_wrapper):
                 super().__init__(app)
@@ -407,10 +406,23 @@ class HttpMcpWrapper:
                 self.logger = logger
 
             async def dispatch(self, request: Request, call_next):
-                # Only handle MCP requests (FastMCP app already only handles /mcp)
-                
-                # Note: Telemetry/tracing now handled in DI function wrapper for unified approach
-                # This middleware focuses purely on session routing
+                # Extract and set trace context from headers for distributed tracing
+                try:
+                    from ..tracing.trace_context_helper import TraceContextHelper
+
+                    # Use helper class for trace context extraction and setup
+                    trace_context = (
+                        await TraceContextHelper.extract_trace_context_from_request(
+                            request
+                        )
+                    )
+                    TraceContextHelper.setup_request_trace_context(
+                        trace_context, self.logger
+                    )
+                except Exception as e:
+                    # Never fail request due to tracing issues
+                    self.logger.warning(f"Failed to set trace context: {e}")
+                    pass
 
                 # Extract session ID from request
                 session_id = await self.http_wrapper._extract_session_id(request)
@@ -443,7 +455,7 @@ class HttpMcpWrapper:
 
         # Add the middleware to FastMCP app
         self._mcp_app.add_middleware(MCPSessionRoutingMiddleware, http_wrapper=self)
-        logger.info("✅ Session routing middleware added to FastMCP app")
+        logger.info("✅ Clean session routing middleware added to FastMCP app (v0.4.0 style)")
 
     async def _extract_session_id(self, request) -> str:
         """Extract session ID from request headers or body."""

_mcp_mesh/engine/unified_mcp_proxy.py

@@ -56,6 +56,67 @@ class UnifiedMCPProxy:
                 f"🔧 UnifiedMCPProxy initialized with kwargs: {self.kwargs_config}"
             )
 
+    def _create_fastmcp_client(self, endpoint: str):
+        """Create FastMCP client with automatic trace header injection.
+
+        This method automatically detects trace context and adds distributed tracing
+        headers when available, while maintaining full backward compatibility.
+
+        Args:
+            endpoint: MCP endpoint URL
+
+        Returns:
+            FastMCP Client instance with or without trace headers
+        """
+        try:
+            from fastmcp import Client
+            from fastmcp.client.transports import StreamableHttpTransport
+
+            # Try to get current trace context for header injection
+            trace_headers = self._get_trace_headers()
+
+            if trace_headers:
+                # Create client with trace headers for distributed tracing
+                transport = StreamableHttpTransport(url=endpoint, headers=trace_headers)
+                return Client(transport)
+            else:
+                # Create standard client when no trace context available
+                return Client(endpoint)
+
+        except ImportError:
+            # If StreamableHttpTransport not available, fall back to standard client
+            from fastmcp import Client
+
+            return Client(endpoint)
+        except Exception as e:
+            # Any other error, fall back to standard client
+            from fastmcp import Client
+
+            return Client(endpoint)
+
+    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
+                }
+                return headers
+            else:
+                return {}
+
+        except Exception as e:
+            # Never fail MCP calls due to tracing issues
+            return {}
+
     def _configure_from_kwargs(self):
         """Auto-configure proxy settings from kwargs."""
         # Basic configuration
@@ -213,74 +274,63 @@ class UnifiedMCPProxy:
         return await self.call_tool_with_tracing(self.function_name, kwargs)
 
     async def call_tool_with_tracing(self, name: str, arguments: dict = None) -> Any:
-        """Call a tool with comprehensive tracing and telemetry."""
-        # Check if telemetry is enabled
-        if not self.telemetry_enabled:
-            self.logger.debug(f"🔇 Telemetry disabled, calling tool directly: {name}")
-            return await self.call_tool(name, arguments)
-
-        # Import tracing dependencies
+        """Call a tool with clean ExecutionTracer integration (v0.4.0 style)."""
+        # Check if telemetry is enabled - use same check as ExecutionTracer for consistency
         from ..tracing.execution_tracer import ExecutionTracer
+        from ..tracing.utils import is_tracing_enabled
 
-        # Initialize execution tracer
-        tracer = ExecutionTracer(name, self.logger)
-
-        try:
-            # Start execution tracing with proxy metadata
-            tracer.start_execution(
-                args=(),
-                kwargs=arguments or {},
-                dependencies=[self.endpoint],
-                mesh_positions=[],
-                injected_count=1,  # This is an injected dependency call
-            )
-
-            # Add proxy-specific metadata
-            tracer.execution_metadata.update(
-                {
-                    "call_type": "unified_mcp_proxy",
-                    "endpoint": self.endpoint,
-                    "function_name": name,
-                    "proxy_type": "fastmcp_with_fallback",
-                    "streaming_capable": self.streaming_capable,
-                    "timeout": self.timeout,
-                    "retry_count": self.retry_count,
-                }
-            )
-
-            # Add enhanced agent context metadata if enabled
-            if self.collect_agent_context:
-                try:
-                    agent_context = self._collect_agent_context_metadata(
-                        name, arguments
-                    )
-                    tracer.execution_metadata.update(agent_context)
-                except Exception as e:
-                    self.logger.debug(f"Failed to collect agent context metadata: {e}")
-                    # Add minimal fallback metadata
-                    tracer.execution_metadata.update(
-                        {
-                            "client_agent_id": "unknown",
-                            "target_agent_endpoint": self.endpoint,
-                            "proxy_instance_id": id(self),
-                        }
-                    )
+        if not self.telemetry_enabled or not is_tracing_enabled():
+            return await self.call_tool(name, arguments)
 
-            # Perform the actual tool call
-            result = await self.call_tool(name, arguments)
+        # Create wrapper function for ExecutionTracer compatibility
+        async def proxy_call_wrapper(*args, **kwargs):
+            # Add proxy-specific metadata to execution context if tracer is available
+            try:
+                from ..tracing.context import TraceContext
+
+                current_trace = TraceContext.get_current()
+                if current_trace and hasattr(current_trace, "execution_metadata"):
+                    # Add proxy metadata to current trace
+                    proxy_metadata = {
+                        "call_type": "unified_mcp_proxy",
+                        "endpoint": self.endpoint,
+                        "proxy_type": "fastmcp_with_fallback",
+                        "streaming_capable": self.streaming_capable,
+                        "timeout": self.timeout,
+                        "retry_count": self.retry_count,
+                    }
 
-            # End tracing with success
-            tracer.end_execution(result, success=True)
+                    # Add enhanced agent context if enabled
+                    if self.collect_agent_context:
+                        try:
+                            agent_context = self._collect_agent_context_metadata(
+                                name, arguments
+                            )
+                            proxy_metadata.update(agent_context)
+                        except Exception as e:
+                            self.logger.debug(
+                                f"Failed to collect agent context metadata: {e}"
+                            )
+
+                    # Update current execution metadata
+                    if hasattr(current_trace, "execution_metadata"):
+                        current_trace.execution_metadata.update(proxy_metadata)
 
-            self.logger.info(f"✅ Traced tool call '{name}' succeeded")
-            return result
+            except Exception as e:
+                self.logger.debug(f"Failed to add proxy metadata: {e}")
 
-        except Exception as e:
-            # End tracing with error
-            tracer.end_execution(error=str(e), success=False)
+            return await self.call_tool(name, arguments)
 
-            self.logger.error(f"❌ Traced tool call '{name}' failed: {e}")
-            raise  # Re-raise the exception
+        # Use ExecutionTracer's static async method for clean integration
+        return await ExecutionTracer.trace_function_execution_async(
+            proxy_call_wrapper,
+            args=(),
+            kwargs={},  # arguments are handled inside the wrapper
+            dependencies=[self.endpoint],
+            mesh_positions=[],
+            injected_count=1,
+            logger_instance=self.logger,
+        )
 
     async def call_tool(self, name: str, arguments: dict = None) -> Any:
         """Call a tool using FastMCP client with HTTP transport.
@@ -292,15 +342,13 @@ class UnifiedMCPProxy:
         start_time = time.time()
 
         try:
-            from fastmcp import Client
-
             # Use correct FastMCP client endpoint - agents expose MCP on /mcp
             mcp_endpoint = f"{self.endpoint}/mcp"
             self.logger.debug(f"🔄 Trying FastMCP client with endpoint: {mcp_endpoint}")
 
-            async with Client(mcp_endpoint) as client:
-                # Add tracing metadata for FastMCP client
-                self.logger.debug(f"🔄 FastMCP client call with tracing: {name}")
+            # Create client with automatic trace header injection
+            client_instance = self._create_fastmcp_client(mcp_endpoint)
+            async with client_instance as client:
 
                 # Use FastMCP's call_tool which returns CallToolResult object
                 result = await client.call_tool(name, arguments or {})
@@ -651,46 +699,51 @@ class UnifiedMCPProxy:
     # MCP Protocol Methods - using FastMCP client's superior implementation
     async def list_tools(self) -> list:
         """List available tools from remote agent."""
-        from fastmcp import Client
-
         mcp_endpoint = f"{self.endpoint}/mcp"
-        async with Client(mcp_endpoint) as client:
+
+        # Create client with automatic trace header injection
+        client_instance = self._create_fastmcp_client(mcp_endpoint)
+        async with client_instance as client:
             result = await client.list_tools()
             return result.tools if hasattr(result, "tools") else result
 
     async def list_resources(self) -> list:
         """List available resources from remote agent."""
-        from fastmcp import Client
-
         mcp_endpoint = f"{self.endpoint}/mcp"
-        async with Client(mcp_endpoint) as client:
+
+        # Create client with automatic trace header injection
+        client_instance = self._create_fastmcp_client(mcp_endpoint)
+        async with client_instance as client:
             result = await client.list_resources()
             return result.resources if hasattr(result, "resources") else result
 
     async def read_resource(self, uri: str) -> Any:
         """Read resource contents from remote agent."""
-        from fastmcp import Client
-
         mcp_endpoint = f"{self.endpoint}/mcp"
-        async with Client(mcp_endpoint) as client:
+
+        # Create client with automatic trace header injection
+        client_instance = self._create_fastmcp_client(mcp_endpoint)
+        async with client_instance as client:
             result = await client.read_resource(uri)
             return result.contents if hasattr(result, "contents") else result
 
     async def list_prompts(self) -> list:
         """List available prompts from remote agent."""
-        from fastmcp import Client
-
         mcp_endpoint = f"{self.endpoint}/mcp"
-        async with Client(mcp_endpoint) as client:
+
+        # Create client with automatic trace header injection
+        client_instance = self._create_fastmcp_client(mcp_endpoint)
+        async with client_instance as client:
             result = await client.list_prompts()
             return result.prompts if hasattr(result, "prompts") else result
 
     async def get_prompt(self, name: str, arguments: dict = None) -> Any:
         """Get prompt template from remote agent."""
-        from fastmcp import Client
-
         mcp_endpoint = f"{self.endpoint}/mcp"
-        async with Client(mcp_endpoint) as client:
+
+        # Create client with automatic trace header injection
+        client_instance = self._create_fastmcp_client(mcp_endpoint)
+        async with client_instance as client:
             result = await client.get_prompt(name, arguments or {})
             return result
 
@@ -700,8 +753,7 @@ class UnifiedMCPProxy:
 
         FastMCP client handles session management internally.
         """
-        import uuid
-        
+
         # Generate session ID for compatibility
         session_id = f"session:{uuid.uuid4().hex[:16]}"
         self.logger.debug(f"📝 Created session ID: {session_id}")

_mcp_mesh/pipeline/api_startup/api_pipeline.py

@@ -11,6 +11,7 @@ import logging
 from ..shared.mesh_pipeline import MeshPipeline
 from .api_server_setup import APIServerSetupStep
 from .fastapi_discovery import FastAPIAppDiscoveryStep
+from .middleware_integration import TracingMiddlewareIntegrationStep
 from .route_collection import RouteCollectionStep
 from .route_integration import RouteIntegrationStep
 
@@ -25,7 +26,8 @@ class APIPipeline(MeshPipeline):
     1. Route collection (@mesh.route decorators)
     2. FastAPI app discovery (find user's FastAPI instances)
     3. Route integration (apply dependency injection)
-    4. API server setup (service registration metadata)
+    4. Tracing middleware integration (add telemetry to FastAPI apps)
+    5. API server setup (service registration metadata)
 
     Unlike MCP agents, API services are consumers so we focus on:
     - Dependency injection into route handlers
@@ -41,10 +43,11 @@ class APIPipeline(MeshPipeline):
         """Setup the API pipeline steps."""
         # Essential API integration steps
         steps = [
-            RouteCollectionStep(),          # Collect @mesh.route decorators
-            FastAPIAppDiscoveryStep(),      # Find user's FastAPI app instances
-            RouteIntegrationStep(),         # Apply dependency injection to routes
-            APIServerSetupStep(),           # Prepare service registration metadata
+            RouteCollectionStep(),              # Collect @mesh.route decorators
+            FastAPIAppDiscoveryStep(),          # Find user's FastAPI app instances
+            RouteIntegrationStep(),             # Apply dependency injection to routes
+            TracingMiddlewareIntegrationStep(), # Add tracing middleware to FastAPI apps
+            APIServerSetupStep(),               # Prepare service registration metadata
             # Note: Heartbeat integration will be added in next phase
             # Note: User controls FastAPI server startup (uvicorn/gunicorn)
         ]