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

_mcp_mesh/__init__.py

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

_mcp_mesh/engine/dependency_injector.py

@@ -14,12 +14,10 @@ import weakref
 from collections.abc import Callable
 from typing import Any
 
-from ..shared.logging_config import (
-    format_log_value,
-    format_result_summary,
-    get_trace_prefix,
-)
-from .signature_analyzer import get_mesh_agent_positions, has_llm_agent_parameter
+from ..shared.logging_config import (format_log_value, format_result_summary,
+                                     get_trace_prefix)
+from .signature_analyzer import (get_mesh_agent_positions,
+                                 has_llm_agent_parameter)
 
 logger = logging.getLogger(__name__)
 

_mcp_mesh/engine/http_wrapper.py

@@ -387,8 +387,13 @@ class HttpMcpWrapper:
                 self.logger = logger
 
             async def dispatch(self, request: Request, call_next):
-                # Extract and set trace context from headers for distributed tracing
+                # Read body once for processing
+                body = await request.body()
+                modified_body = body
+
+                # Extract and set trace context from headers and arguments
                 try:
+                    from ..tracing.context import TraceContext
                     from ..tracing.trace_context_helper import TraceContextHelper
 
                     # DEBUG: Log incoming headers for trace propagation debugging
@@ -399,22 +404,67 @@ class HttpMcpWrapper:
                         f"X-Parent-Span={parent_span_header}, path={request.url.path}"
                     )
 
-                    # Use helper class for trace context extraction and setup
-                    trace_context = (
-                        await TraceContextHelper.extract_trace_context_from_request(
-                            request
+                    # Extract trace context from both headers AND arguments
+                    trace_id = trace_id_header
+                    parent_span = parent_span_header
+
+                    # Try extracting from JSON-RPC body arguments as fallback
+                    # Also strip trace fields from arguments to avoid Pydantic validation errors
+                    if body:
+                        try:
+                            payload = json.loads(body.decode("utf-8"))
+                            if payload.get("method") == "tools/call":
+                                arguments = payload.get("params", {}).get(
+                                    "arguments", {}
+                                )
+
+                                # Extract trace context from arguments (TypeScript uses _trace_id/_parent_span)
+                                if not trace_id and arguments.get("_trace_id"):
+                                    trace_id = arguments.get("_trace_id")
+                                if not parent_span and arguments.get("_parent_span"):
+                                    parent_span = arguments.get("_parent_span")
+
+                                # Strip trace context fields from arguments before passing to FastMCP
+                                if (
+                                    "_trace_id" in arguments
+                                    or "_parent_span" in arguments
+                                ):
+                                    arguments.pop("_trace_id", None)
+                                    arguments.pop("_parent_span", None)
+                                    # Update payload with cleaned arguments
+                                    modified_body = json.dumps(payload).encode("utf-8")
+                                    self.logger.debug(
+                                        f"🔗 Stripped trace fields from arguments, "
+                                        f"trace_id={trace_id[:8] if trace_id else None}..."
+                                    )
+                        except Exception as e:
+                            self.logger.debug(
+                                f"Failed to process body for trace context: {e}"
+                            )
+
+                    # Setup trace context if we have a trace_id
+                    if trace_id:
+                        trace_context = {
+                            "trace_id": trace_id,
+                            "parent_span": parent_span,
+                        }
+                        TraceContextHelper.setup_request_trace_context(
+                            trace_context, self.logger
                         )
-                    )
-                    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
 
+                # Create a new request scope with the modified body
+                async def receive():
+                    return {"type": "http.request", "body": modified_body}
+
+                # Update request with modified receive
+                request._receive = receive
+
                 # Extract session ID from request
-                session_id = await self.http_wrapper._extract_session_id(request)
+                session_id = await self.http_wrapper._extract_session_id_from_body(body)
 
                 if session_id:
                     # Check for existing session assignment
@@ -458,6 +508,15 @@ class HttpMcpWrapper:
         # Try extracting from JSON-RPC body
         try:
             body = await request.body()
+            return await self._extract_session_id_from_body(body)
+        except Exception:
+            pass
+
+        return None
+
+    async def _extract_session_id_from_body(self, body: bytes) -> str:
+        """Extract session ID from already-read request body."""
+        try:
             if body:
                 payload = json.loads(body.decode("utf-8"))
                 if payload.get("method") == "tools/call":

_mcp_mesh/engine/mesh_llm_agent.py

@@ -14,12 +14,8 @@ from typing import Any, Dict, List, Literal, Optional, Union
 from pydantic import BaseModel
 
 from .llm_config import LLMConfig
-from .llm_errors import (
-    LLMAPIError,
-    MaxIterationsError,
-    ResponseParseError,
-    ToolExecutionError,
-)
+from .llm_errors import (LLMAPIError, MaxIterationsError, ResponseParseError,
+                         ToolExecutionError)
 from .provider_handlers import ProviderHandlerRegistry
 from .response_parser import ResponseParser
 from .tool_executor import ToolExecutor
@@ -27,7 +23,8 @@ from .tool_schema_builder import ToolSchemaBuilder
 
 # Import Jinja2 for template rendering
 try:
-    from jinja2 import Environment, FileSystemLoader, Template, TemplateSyntaxError
+    from jinja2 import (Environment, FileSystemLoader, Template,
+                        TemplateSyntaxError)
 except ImportError:
     Environment = None
     FileSystemLoader = None

_mcp_mesh/engine/mesh_llm_agent_injector.py

@@ -431,7 +431,8 @@ class MeshLlmAgentInjector(BaseInjector):
                     if is_template:
                         # Templates enabled - create per-call agent with context
                         # Import signature analyzer for context detection
-                        from .signature_analyzer import get_context_parameter_name
+                        from .signature_analyzer import \
+                            get_context_parameter_name
 
                         # Detect context parameter
                         context_param_name = config_dict.get("context_param")

_mcp_mesh/engine/provider_handlers/__init__.py

@@ -5,15 +5,28 @@ This package provides vendor-specific customization for different LLM providers
 (Claude, OpenAI, Gemini, etc.) to optimize API calls and response handling.
 """
 
-from .base_provider_handler import BaseProviderHandler
+from .base_provider_handler import (
+    BASE_TOOL_INSTRUCTIONS,
+    CLAUDE_ANTI_XML_INSTRUCTION,
+    BaseProviderHandler,
+    make_schema_strict,
+)
 from .claude_handler import ClaudeHandler
+from .gemini_handler import GeminiHandler
 from .generic_handler import GenericHandler
 from .openai_handler import OpenAIHandler
 from .provider_handler_registry import ProviderHandlerRegistry
 
 __all__ = [
+    # Constants
+    "BASE_TOOL_INSTRUCTIONS",
+    "CLAUDE_ANTI_XML_INSTRUCTION",
+    # Utilities
+    "make_schema_strict",
+    # Handlers
     "BaseProviderHandler",
     "ClaudeHandler",
+    "GeminiHandler",
     "OpenAIHandler",
     "GenericHandler",
     "ProviderHandlerRegistry",

_mcp_mesh/engine/provider_handlers/base_provider_handler.py

@@ -5,11 +5,117 @@ This module defines the abstract base class for provider-specific handlers
 that customize how different LLM vendors (Claude, OpenAI, Gemini, etc.) are called.
 """
 
+import copy
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional
+from typing import Any, Optional
 
 from pydantic import BaseModel
 
+# ============================================================================
+# Shared Constants
+# ============================================================================
+
+# Base tool calling instructions shared across all providers.
+# Claude handler adds anti-XML instruction on top of this.
+BASE_TOOL_INSTRUCTIONS = """
+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
+"""
+
+# Anti-XML instruction for Claude (prevents <invoke> style tool calls).
+CLAUDE_ANTI_XML_INSTRUCTION = (
+    '- NEVER use XML-style syntax like <invoke name="tool_name"/>'
+)
+
+
+# ============================================================================
+# Shared Schema Utilities
+# ============================================================================
+
+
+def make_schema_strict(
+    schema: dict[str, Any],
+    add_all_required: bool = True,
+) -> dict[str, Any]:
+    """
+    Make a JSON schema strict for structured output.
+
+    This is a shared utility used by OpenAI, Gemini, and Claude handlers.
+    Adds additionalProperties: false to all object types and optionally
+    ensures 'required' includes all property keys.
+
+    Args:
+        schema: JSON schema to make strict
+        add_all_required: If True, set 'required' to include ALL property keys.
+                         OpenAI and Gemini require this; Claude does not.
+                         Default: True
+
+    Returns:
+        New schema with strict constraints (original not mutated)
+    """
+    result = copy.deepcopy(schema)
+    _add_strict_constraints_recursive(result, add_all_required)
+    return result
+
+
+def _add_strict_constraints_recursive(obj: Any, add_all_required: bool) -> None:
+    """
+    Recursively add strict constraints to a schema object.
+
+    Args:
+        obj: Schema object to process (mutated in place)
+        add_all_required: Whether to set required to all property keys
+    """
+    if not isinstance(obj, dict):
+        return
+
+    # If this is an object type, add additionalProperties: false
+    if obj.get("type") == "object":
+        obj["additionalProperties"] = False
+
+        # Optionally set required to include all property keys
+        if add_all_required and "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():
+            _add_strict_constraints_recursive(def_schema, add_all_required)
+
+    # Process properties
+    if "properties" in obj:
+        for prop_schema in obj["properties"].values():
+            _add_strict_constraints_recursive(prop_schema, add_all_required)
+
+    # Process items (for arrays)
+    # items can be an object (single schema) or a list (tuple validation in older drafts)
+    if "items" in obj:
+        items = obj["items"]
+        if isinstance(items, dict):
+            _add_strict_constraints_recursive(items, add_all_required)
+        elif isinstance(items, list):
+            for item in items:
+                _add_strict_constraints_recursive(item, add_all_required)
+
+    # Process prefixItems (tuple validation in JSON Schema draft 2020-12)
+    if "prefixItems" in obj:
+        for item in obj["prefixItems"]:
+            _add_strict_constraints_recursive(item, add_all_required)
+
+    # Process anyOf, oneOf, allOf
+    for key in ("anyOf", "oneOf", "allOf"):
+        if key in obj:
+            for item in obj[key]:
+                _add_strict_constraints_recursive(item, add_all_required)
+
+
+# ============================================================================
+# Base Provider Handler
+# ============================================================================
+
 
 class BaseProviderHandler(ABC):
     """
@@ -43,11 +149,11 @@ class BaseProviderHandler(ABC):
     @abstractmethod
     def prepare_request(
         self,
-        messages: List[Dict[str, Any]],
-        tools: Optional[List[Dict[str, Any]]],
+        messages: list[dict[str, Any]],
+        tools: Optional[list[dict[str, Any]]],
         output_type: type[BaseModel],
-        **kwargs: Any
-    ) -> Dict[str, Any]:
+        **kwargs: Any,
+    ) -> dict[str, Any]:
         """
         Prepare vendor-specific request parameters.
 
@@ -74,8 +180,8 @@ class BaseProviderHandler(ABC):
     def format_system_prompt(
         self,
         base_prompt: str,
-        tool_schemas: Optional[List[Dict[str, Any]]],
-        output_type: type[BaseModel]
+        tool_schemas: Optional[list[dict[str, Any]]],
+        output_type: type[BaseModel],
     ) -> str:
         """
         Format system prompt for vendor-specific requirements.
@@ -95,7 +201,7 @@ class BaseProviderHandler(ABC):
         """
         pass
 
-    def get_vendor_capabilities(self) -> Dict[str, bool]:
+    def get_vendor_capabilities(self) -> dict[str, bool]:
         """
         Return vendor-specific capability flags.
 

_mcp_mesh/engine/provider_handlers/claude_handler.py

@@ -17,11 +17,16 @@ Features:
 
 import json
 import logging
-from typing import Any, Optional, get_args, get_origin
+from typing import Any, Optional
 
 from pydantic import BaseModel
 
-from .base_provider_handler import BaseProviderHandler
+from .base_provider_handler import (
+    BASE_TOOL_INSTRUCTIONS,
+    BaseProviderHandler,
+    CLAUDE_ANTI_XML_INSTRUCTION,
+    make_schema_strict,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -141,51 +146,6 @@ class ClaudeHandler(BaseProviderHandler):
         # Default to strict for unknown types
         return OUTPUT_MODE_STRICT
 
-    def _make_schema_strict(self, schema: dict[str, Any]) -> dict[str, Any]:
-        """
-        Make a JSON schema strict for Claude's structured output.
-
-        Claude requires additionalProperties: false on all object types.
-        This recursively processes the schema to add this constraint.
-
-        Args:
-            schema: JSON schema dict
-
-        Returns:
-            Schema with additionalProperties: false on all objects
-        """
-        if not isinstance(schema, dict):
-            return schema
-
-        result = schema.copy()
-
-        # If this is an object type, add additionalProperties: false
-        if result.get("type") == "object":
-            result["additionalProperties"] = False
-
-        # Recursively process nested schemas
-        if "properties" in result:
-            result["properties"] = {
-                k: self._make_schema_strict(v) for k, v in result["properties"].items()
-            }
-
-        # Process $defs (Pydantic uses this for nested models)
-        if "$defs" in result:
-            result["$defs"] = {
-                k: self._make_schema_strict(v) for k, v in result["$defs"].items()
-            }
-
-        # Process items for arrays
-        if "items" in result:
-            result["items"] = self._make_schema_strict(result["items"])
-
-        # Process anyOf, oneOf, allOf
-        for key in ["anyOf", "oneOf", "allOf"]:
-            if key in result:
-                result[key] = [self._make_schema_strict(s) for s in result[key]]
-
-        return result
-
     def _apply_prompt_caching(
         self, messages: list[dict[str, Any]]
     ) -> list[dict[str, Any]]:
@@ -302,9 +262,10 @@ class ClaudeHandler(BaseProviderHandler):
         # Only add response_format in "strict" mode
         if determined_mode == OUTPUT_MODE_STRICT:
             # Claude requires additionalProperties: false on all object types
+            # Unlike OpenAI/Gemini, Claude doesn't require all properties in 'required'
             if isinstance(output_type, type) and issubclass(output_type, BaseModel):
                 schema = output_type.model_json_schema()
-                strict_schema = self._make_schema_strict(schema)
+                strict_schema = make_schema_strict(schema, add_all_required=False)
                 request_params["response_format"] = {
                     "type": "json_schema",
                     "json_schema": {
@@ -346,15 +307,12 @@ class ClaudeHandler(BaseProviderHandler):
         # Add tool calling instructions if tools available
         # These prevent Claude from using XML-style <invoke> syntax
         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
-- NEVER use XML-style syntax like <invoke name="tool_name"/>
-- After receiving tool results, you can make additional calls if needed
-- Once you have all needed information, provide your final response
-"""
+            # Use base instructions but insert anti-XML rule for Claude
+            instructions = BASE_TOOL_INSTRUCTIONS.replace(
+                "- Make ONE tool call at a time",
+                f"- Make ONE tool call at a time\n{CLAUDE_ANTI_XML_INSTRUCTION}",
+            )
+            system_content += instructions
 
         # Add output format instructions based on mode
         if determined_mode == OUTPUT_MODE_TEXT:

_mcp_mesh/engine/provider_handlers/gemini_handler.py

@@ -0,0 +1,181 @@
+"""
+Gemini/Google provider handler.
+
+Optimized for Gemini models (Gemini 2.0 Flash, Gemini 1.5 Pro, etc.)
+using Google's best practices for tool calling and structured output.
+
+Features:
+- Native structured output via response_format (similar to OpenAI)
+- Native function calling support
+- Support for Gemini 2.x and 3.x models
+- Large context windows (up to 2M tokens)
+
+Reference:
+- https://docs.litellm.ai/docs/providers/gemini
+- https://ai.google.dev/gemini-api/docs
+"""
+
+import logging
+from typing import Any, Optional
+
+from pydantic import BaseModel
+
+from .base_provider_handler import (
+    BASE_TOOL_INSTRUCTIONS,
+    BaseProviderHandler,
+    make_schema_strict,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class GeminiHandler(BaseProviderHandler):
+    """
+    Provider handler for Google Gemini models.
+
+    Gemini Characteristics:
+    - Native structured output via response_format parameter (LiteLLM translates)
+    - Native function calling support
+    - Large context windows (1M-2M tokens)
+    - Multimodal support (text, images, video, audio)
+    - Works well with concise, focused prompts
+
+    Key Similarities with OpenAI:
+    - Uses response_format for structured output (via LiteLLM translation)
+    - Native function calling format
+    - Similar schema enforcement requirements
+
+    Supported Models (via LiteLLM):
+    - gemini/gemini-2.0-flash (fast, efficient)
+    - gemini/gemini-2.0-flash-lite (fastest, most efficient)
+    - gemini/gemini-1.5-pro (high capability)
+    - gemini/gemini-1.5-flash (balanced)
+    - gemini/gemini-3-flash-preview (reasoning support)
+    - gemini/gemini-3-pro-preview (advanced reasoning)
+
+    Reference:
+    https://docs.litellm.ai/docs/providers/gemini
+    """
+
+    def __init__(self):
+        """Initialize Gemini handler."""
+        super().__init__(vendor="gemini")
+
+    def prepare_request(
+        self,
+        messages: list[dict[str, Any]],
+        tools: Optional[list[dict[str, Any]]],
+        output_type: type,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        """
+        Prepare request parameters for Gemini API via LiteLLM.
+
+        Gemini Strategy:
+        - Use response_format parameter for structured JSON output
+        - LiteLLM handles translation to Gemini's native format
+        - Skip structured output for str return types (text mode)
+
+        Args:
+            messages: List of message dicts
+            tools: Optional list of tool schemas
+            output_type: Return type (str or Pydantic model)
+            **kwargs: Additional model parameters
+
+        Returns:
+            Dictionary of parameters for litellm.completion()
+        """
+        # Build base request
+        request_params = {
+            "messages": messages,
+            **kwargs,  # Pass through temperature, max_tokens, etc.
+        }
+
+        # Add tools if provided
+        # LiteLLM will convert OpenAI tool format to Gemini's function_declarations
+        if tools:
+            request_params["tools"] = tools
+
+        # Skip structured output for str return type (text mode)
+        if output_type is str:
+            return request_params
+
+        # Only add response_format for Pydantic models
+        if not (isinstance(output_type, type) and issubclass(output_type, BaseModel)):
+            return request_params
+
+        # Add response_format for structured output
+        # LiteLLM translates this to Gemini's native format
+        schema = output_type.model_json_schema()
+
+        # Transform schema for strict mode compliance
+        # Gemini requires additionalProperties: false and all properties in required
+        schema = make_schema_strict(schema, add_all_required=True)
+
+        # Gemini structured output format (via LiteLLM)
+        request_params["response_format"] = {
+            "type": "json_schema",
+            "json_schema": {
+                "name": output_type.__name__,
+                "schema": schema,
+                "strict": True,  # Enforce schema compliance
+            },
+        }
+
+        return request_params
+
+    def format_system_prompt(
+        self,
+        base_prompt: str,
+        tool_schemas: Optional[list[dict[str, Any]]],
+        output_type: type,
+    ) -> str:
+        """
+        Format system prompt for Gemini (concise approach).
+
+        Gemini Strategy:
+        1. Use base prompt as-is
+        2. Add tool calling instructions if tools present
+        3. Minimal JSON instructions (response_format handles structure)
+        4. Keep prompt concise - Gemini works well with clear, direct prompts
+
+        Args:
+            base_prompt: Base system prompt
+            tool_schemas: Optional tool schemas
+            output_type: Expected response type (str or Pydantic model)
+
+        Returns:
+            Formatted system prompt optimized for Gemini
+        """
+        system_content = base_prompt
+
+        # Add tool calling instructions if tools available
+        if tool_schemas:
+            system_content += BASE_TOOL_INSTRUCTIONS
+
+        # Skip JSON note for str return type (text mode)
+        if output_type is str:
+            return system_content
+
+        # Add brief JSON note (response_format handles enforcement)
+        if isinstance(output_type, type) and issubclass(output_type, BaseModel):
+            system_content += f"\n\nYour final response will be structured as JSON matching the {output_type.__name__} format."
+
+        return system_content
+
+    def get_vendor_capabilities(self) -> dict[str, bool]:
+        """
+        Return Gemini-specific capabilities.
+
+        Returns:
+            Capability flags for Gemini
+        """
+        return {
+            "native_tool_calling": True,  # Gemini has native function calling
+            "structured_output": True,  # Supports structured output via response_format
+            "streaming": True,  # Supports streaming
+            "vision": True,  # Gemini supports multimodal (images, video, audio)
+            "json_mode": True,  # Native JSON mode via response_format
+            "large_context": True,  # Up to 2M tokens context window
+        }
+