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

_mcp_mesh/engine/provider_handlers/openai_handler.py

@@ -5,12 +5,15 @@ Optimized for OpenAI models (GPT-4, GPT-4 Turbo, GPT-3.5-turbo)
 using OpenAI's native structured output capabilities.
 """
 
-import json
 from typing import Any, Optional
 
 from pydantic import BaseModel
 
-from .base_provider_handler import BaseProviderHandler
+from .base_provider_handler import (
+    BASE_TOOL_INSTRUCTIONS,
+    BaseProviderHandler,
+    make_schema_strict,
+)
 
 
 class OpenAIHandler(BaseProviderHandler):
@@ -94,8 +97,8 @@ class OpenAIHandler(BaseProviderHandler):
             schema = output_type.model_json_schema()
 
             # Transform schema for OpenAI strict mode
-            # OpenAI requires additionalProperties: false on all object schemas
-            schema = self._add_additional_properties_false(schema)
+            # OpenAI requires additionalProperties: false and all properties in required
+            schema = make_schema_strict(schema, add_all_required=True)
 
             # OpenAI structured output format
             # See: https://platform.openai.com/docs/guides/structured-outputs
@@ -143,14 +146,7 @@ class OpenAIHandler(BaseProviderHandler):
 
         # Add tool calling instructions if tools available
         if tool_schemas:
-            system_content += """
-
-IMPORTANT TOOL CALLING RULES:
-- You have access to tools that you can call to gather information
-- Make ONE tool call at a time
-- After receiving tool results, you can make additional calls if needed
-- Once you have all needed information, provide your final response
-"""
+            system_content += BASE_TOOL_INSTRUCTIONS
 
         # Skip JSON note for str return type (text mode)
         if output_type is str:
@@ -181,54 +177,3 @@ IMPORTANT TOOL CALLING RULES:
             "json_mode": True,  # Has dedicated JSON mode via response_format
         }
 
-    def _add_additional_properties_false(
-        self, schema: dict[str, Any]
-    ) -> dict[str, Any]:
-        """
-        Recursively add additionalProperties: false to all object schemas.
-
-        OpenAI strict mode requires this for all object schemas.
-        See: https://platform.openai.com/docs/guides/structured-outputs
-
-        Args:
-            schema: JSON schema from Pydantic model
-
-        Returns:
-            Modified schema with additionalProperties: false on all objects
-        """
-        import copy
-
-        schema = copy.deepcopy(schema)
-        self._add_additional_properties_recursive(schema)
-        return schema
-
-    def _add_additional_properties_recursive(self, obj: Any) -> None:
-        """Recursively process schema for OpenAI strict mode compliance."""
-        if isinstance(obj, dict):
-            # If this is an object type, add additionalProperties: false
-            # and ensure required includes all properties
-            if obj.get("type") == "object":
-                obj["additionalProperties"] = False
-                # OpenAI strict mode: required must include ALL property keys
-                if "properties" in obj:
-                    obj["required"] = list(obj["properties"].keys())
-
-            # Process $defs (Pydantic uses this for nested models)
-            if "$defs" in obj:
-                for def_schema in obj["$defs"].values():
-                    self._add_additional_properties_recursive(def_schema)
-
-            # Process properties
-            if "properties" in obj:
-                for prop_schema in obj["properties"].values():
-                    self._add_additional_properties_recursive(prop_schema)
-
-            # Process items (for arrays)
-            if "items" in obj:
-                self._add_additional_properties_recursive(obj["items"])
-
-            # Process anyOf, oneOf, allOf
-            for key in ("anyOf", "oneOf", "allOf"):
-                if key in obj:
-                    for item in obj[key]:
-                        self._add_additional_properties_recursive(item)

_mcp_mesh/engine/provider_handlers/provider_handler_registry.py

@@ -5,10 +5,11 @@ Manages selection and instantiation of provider handlers based on vendor name.
 """
 
 import logging
-from typing import Dict, Optional, Type
+from typing import Optional
 
 from .base_provider_handler import BaseProviderHandler
 from .claude_handler import ClaudeHandler
+from .gemini_handler import GeminiHandler
 from .generic_handler import GenericHandler
 from .openai_handler import OpenAIHandler
 
@@ -39,16 +40,17 @@ class ProviderHandlerRegistry:
     """
 
     # Built-in vendor mappings
-    _handlers: Dict[str, Type[BaseProviderHandler]] = {
+    _handlers: dict[str, type[BaseProviderHandler]] = {
         "anthropic": ClaudeHandler,
         "openai": OpenAIHandler,
+        "gemini": GeminiHandler,
     }
 
     # Cache of instantiated handlers (singleton per vendor)
-    _instances: Dict[str, BaseProviderHandler] = {}
+    _instances: dict[str, BaseProviderHandler] = {}
 
     @classmethod
-    def register(cls, vendor: str, handler_class: Type[BaseProviderHandler]) -> None:
+    def register(cls, vendor: str, handler_class: type[BaseProviderHandler]) -> None:
         """
         Register a custom provider handler.
 
@@ -73,7 +75,9 @@ class ProviderHandlerRegistry:
             )
 
         cls._handlers[vendor] = handler_class
-        logger.info(f"📝 Registered provider handler: {vendor} -> {handler_class.__name__}")
+        logger.info(
+            f"📝 Registered provider handler: {vendor} -> {handler_class.__name__}"
+        )
 
         # Clear cached instance if it exists (force re-instantiation)
         if vendor in cls._instances:
@@ -119,9 +123,7 @@ class ProviderHandlerRegistry:
         # Get handler class (or fallback to Generic)
         if vendor in cls._handlers:
             handler_class = cls._handlers[vendor]
-            logger.info(
-                f"✅ Selected {handler_class.__name__} for vendor: {vendor}"
-            )
+            logger.info(f"✅ Selected {handler_class.__name__} for vendor: {vendor}")
         else:
             handler_class = GenericHandler
             if vendor != "unknown":
@@ -132,14 +134,18 @@ class ProviderHandlerRegistry:
                 logger.debug("Using GenericHandler for unknown vendor")
 
         # Instantiate and cache
-        handler = handler_class() if handler_class != GenericHandler else GenericHandler(vendor)
+        handler = (
+            handler_class()
+            if handler_class != GenericHandler
+            else GenericHandler(vendor)
+        )
         cls._instances[vendor] = handler
 
         logger.debug(f"🆕 Instantiated handler: {handler}")
         return handler
 
     @classmethod
-    def list_vendors(cls) -> Dict[str, str]:
+    def list_vendors(cls) -> dict[str, str]:
         """
         List all registered vendors and their handlers.
 

_mcp_mesh/engine/response_parser.py

@@ -94,8 +94,9 @@ class ResponseParser:
 
         Tries multiple strategies to find JSON in mixed responses:
         1. Find ```json ... ``` code fence blocks
-        2. Find any JSON object {...} in the content
-        3. Return original content if no extraction needed
+        2. Find any JSON object {...} using progressive json.loads
+        3. Find any JSON array [...] using progressive json.loads
+        4. Return original content if no extraction needed
 
         Args:
             content: Raw content that may contain narrative, XML, and JSON
@@ -109,25 +110,70 @@ class ResponseParser:
             extracted = json_match.group(1).strip()
             return extracted
 
-        # Strategy 2: Try to find any JSON object {...} in content
-        # Look for balanced braces starting with { and ending with }
+        # Strategy 2: Try to find JSON object using progressive json.loads
+        # This correctly handles braces inside string values
         brace_start = content.find("{")
         if brace_start != -1:
-            # Find matching closing brace
-            brace_count = 0
-            for i in range(brace_start, len(content)):
-                if content[i] == "{":
-                    brace_count += 1
-                elif content[i] == "}":
-                    brace_count -= 1
-                    if brace_count == 0:
-                        # Found matching brace
-                        extracted = content[brace_start : i + 1]
-                        return extracted
+            result = ResponseParser._try_progressive_parse(
+                content, brace_start, "{", "}"
+            )
+            if result:
+                return result
+
+        # Strategy 3: Try to find JSON array using progressive json.loads
+        bracket_start = content.find("[")
+        if bracket_start != -1:
+            result = ResponseParser._try_progressive_parse(
+                content, bracket_start, "[", "]"
+            )
+            if result:
+                return result
 
         # No JSON found, return original
         return content
 
+    @staticmethod
+    def _try_progressive_parse(
+        content: str, start: int, open_char: str, close_char: str
+    ) -> str | None:
+        """
+        Try to extract valid JSON by progressively extending the end position.
+        This correctly handles braces/brackets inside string values.
+
+        Args:
+            content: The full content string
+            start: Starting index of the JSON
+            open_char: Opening character ('{' or '[')
+            close_char: Closing character ('}' or ']')
+
+        Returns:
+            Extracted JSON string or None if not found
+        """
+        # Find potential end positions based on depth counting
+        depth = 0
+        potential_ends: list[int] = []
+
+        for i in range(start, len(content)):
+            char = content[i]
+            if char == open_char:
+                depth += 1
+            elif char == close_char:
+                depth -= 1
+                if depth == 0:
+                    potential_ends.append(i)
+
+        # Try each potential end position (shortest first for efficiency)
+        for end in potential_ends:
+            candidate = content[start : end + 1]
+            try:
+                json.loads(candidate)
+                return candidate
+            except json.JSONDecodeError:
+                # Not valid JSON, try next potential end
+                continue
+
+        return None
+
     @staticmethod
     def _strip_markdown_fences(content: str) -> str:
         """

_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__)
 
@@ -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