mcp-mesh 0.7.12__py3-none-any.whl → 0.7.13__py3-none-any.whl

_mcp_mesh/engine/mesh_llm_agent.py

@@ -7,6 +7,7 @@ Provides automatic agentic loop for LLM-based agents with tool integration.
 import asyncio
 import json
 import logging
+import time
 from pathlib import Path
 from typing import Any, Dict, List, Literal, Optional, Union
 
@@ -67,6 +68,7 @@ class MeshLlmAgent:
         context_value: Optional[Any] = None,
         provider_proxy: Optional[Any] = None,
         vendor: Optional[str] = None,
+        default_model_params: Optional[dict[str, Any]] = None,
     ):
         """
         Initialize MeshLlmAgent proxy.
@@ -80,6 +82,9 @@ class MeshLlmAgent:
             context_value: Optional context for template rendering (MeshContextModel, dict, or None)
             provider_proxy: Optional pre-resolved provider proxy for mesh delegation
             vendor: Optional vendor name for handler selection (e.g., "anthropic", "openai")
+            default_model_params: Optional dict of default LLM parameters from decorator
+                                  (e.g., max_tokens, temperature). These are merged with
+                                  call-time kwargs, with call-time taking precedence.
         """
         self.config = config
         self.provider = config.provider
@@ -92,6 +97,9 @@ class MeshLlmAgent:
         self.system_prompt = config.system_prompt  # Public attribute for tests
         self.output_mode = config.output_mode  # Output mode override (strict/hint/text)
         self._iteration_count = 0
+        self._default_model_params = (
+            default_model_params or {}
+        )  # Decorator-level defaults
 
         # Detect if using mesh delegation (provider is dict)
         self._is_mesh_delegated = isinstance(self.provider, dict)
@@ -336,6 +344,64 @@ IMPORTANT TOOL CALLING RULES:
         # Otherwise, use literal system prompt from config
         return self.system_prompt or ""
 
+    def _attach_mesh_meta(
+        self,
+        result: Any,
+        model: str,
+        input_tokens: int,
+        output_tokens: int,
+        latency_ms: float,
+    ) -> Any:
+        """
+        Attach _mesh_meta to result object if possible.
+
+        For Pydantic models and regular classes, attaches LlmMeta as _mesh_meta.
+        For primitives (str, int, etc.) and frozen models, silently skips.
+
+        Args:
+            result: The parsed result object
+            model: Model identifier used
+            input_tokens: Total input tokens across all iterations
+            output_tokens: Total output tokens across all iterations
+            latency_ms: Total latency in milliseconds
+
+        Returns:
+            The result object (unchanged, but with _mesh_meta attached if possible)
+        """
+        from mesh.types import LlmMeta
+
+        # Extract provider from model string (e.g., "anthropic/claude-3-5-haiku" -> "anthropic")
+        provider = "unknown"
+        if isinstance(model, str) and "/" in model:
+            provider = model.split("/")[0]
+
+        meta = LlmMeta(
+            provider=provider,
+            model=model or "unknown",
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            total_tokens=input_tokens + output_tokens,
+            latency_ms=latency_ms,
+        )
+
+        # Try to attach _mesh_meta to result
+        try:
+            # This works for Pydantic models and most Python objects
+            object.__setattr__(result, "_mesh_meta", meta)
+            logger.debug(
+                f"📊 Attached _mesh_meta: model={model}, "
+                f"tokens={input_tokens}+{output_tokens}={input_tokens + output_tokens}, "
+                f"latency={latency_ms:.1f}ms"
+            )
+        except (TypeError, AttributeError):
+            # Primitives (str, int, etc.) and frozen objects don't support attribute assignment
+            logger.debug(
+                f"📊 Could not attach _mesh_meta to {type(result).__name__} "
+                f"(tokens={input_tokens}+{output_tokens}, latency={latency_ms:.1f}ms)"
+            )
+
+        return result
+
     async def _get_mesh_provider(self) -> Any:
         """
         Get the mesh provider proxy (already resolved during heartbeat).
@@ -456,9 +522,20 @@ IMPORTANT TOOL CALLING RULES:
                     self.message = message
                     self.finish_reason = "stop"
 
+            # Issue #311: Mock usage object for token tracking
+            class MockUsage:
+                def __init__(self, usage_dict):
+                    self.prompt_tokens = usage_dict.get("prompt_tokens", 0)
+                    self.completion_tokens = usage_dict.get("completion_tokens", 0)
+                    self.total_tokens = self.prompt_tokens + self.completion_tokens
+
             class MockResponse:
                 def __init__(self, message_dict):
                     self.choices = [MockChoice(MockMessage(message_dict))]
+                    # Issue #311: Extract usage from _mesh_usage if present
+                    mesh_usage = message_dict.get("_mesh_usage")
+                    self.usage = MockUsage(mesh_usage) if mesh_usage else None
+                    self.model = mesh_usage.get("model") if mesh_usage else None
 
             logger.debug(
                 f"📥 Received response from mesh provider: "
@@ -523,6 +600,12 @@ IMPORTANT TOOL CALLING RULES:
         """
         self._iteration_count = 0
 
+        # Issue #311: Track timing and token usage for _mesh_meta
+        start_time = time.perf_counter()
+        total_input_tokens = 0
+        total_output_tokens = 0
+        effective_model = self.model  # Track actual model used
+
         # Check if litellm is available
         if completion is None:
             raise ImportError(
@@ -583,11 +666,15 @@ IMPORTANT TOOL CALLING RULES:
             try:
                 # Call LLM (either direct LiteLLM or mesh-delegated)
                 try:
+                    # Merge decorator-level defaults with call-time kwargs
+                    # Call-time kwargs take precedence over defaults
+                    effective_kwargs = {**self._default_model_params, **kwargs}
+
                     # Build kwargs with output_mode override if set
                     call_kwargs = (
-                        {**kwargs, "output_mode": self.output_mode}
+                        {**effective_kwargs, "output_mode": self.output_mode}
                         if self.output_mode
-                        else kwargs
+                        else effective_kwargs
                     )
 
                     # Use provider handler to prepare vendor-specific request
@@ -600,7 +687,7 @@ IMPORTANT TOOL CALLING RULES:
 
                     if self._is_mesh_delegated:
                         # Mesh delegation: extract model_params to send to provider
-                        # Exclude messages/tools (separate params), model/api_key (provider has them),
+                        # Exclude messages/tools (separate params), api_key (provider has it),
                         # and output_mode (only used locally by prepare_request)
                         model_params = {
                             k: v
@@ -609,12 +696,18 @@ IMPORTANT TOOL CALLING RULES:
                             not in [
                                 "messages",
                                 "tools",
-                                "model",
                                 "api_key",
                                 "output_mode",
+                                "model",  # Model handled separately below
                             ]
                         }
 
+                        # Issue #308: Include model override if explicitly set by consumer
+                        # This allows consumer to request a specific model from the provider
+                        # (e.g., use haiku instead of provider's default sonnet)
+                        if self.model:
+                            model_params["model"] = self.model
+
                         logger.debug(
                             f"📤 Delegating to mesh provider with handler-prepared params: "
                             f"keys={list(model_params.keys())}"
@@ -645,6 +738,16 @@ IMPORTANT TOOL CALLING RULES:
                         original_error=e,
                     ) from e
 
+                # Issue #311: Extract token usage from response
+                if hasattr(response, "usage") and response.usage:
+                    usage = response.usage
+                    total_input_tokens += getattr(usage, "prompt_tokens", 0) or 0
+                    total_output_tokens += getattr(usage, "completion_tokens", 0) or 0
+
+                # Issue #311: Track effective model (may differ from requested in mesh delegation)
+                if hasattr(response, "model") and response.model:
+                    effective_model = response.model
+
                 # Extract response content
                 assistant_message = response.choices[0].message
 
@@ -675,7 +778,18 @@ IMPORTANT TOOL CALLING RULES:
                     f"📥 Raw LLM response: {assistant_message.content[:500]}..."
                 )
 
-                return self._parse_response(assistant_message.content)
+                # Parse the response
+                result = self._parse_response(assistant_message.content)
+
+                # Issue #311: Calculate latency and attach _mesh_meta
+                latency_ms = (time.perf_counter() - start_time) * 1000
+                return self._attach_mesh_meta(
+                    result=result,
+                    model=effective_model,
+                    input_tokens=total_input_tokens,
+                    output_tokens=total_output_tokens,
+                    latency_ms=latency_ms,
+                )
 
             except LLMAPIError:
                 # Re-raise LLM API errors as-is

_mcp_mesh/engine/mesh_llm_agent_injector.py

@@ -537,6 +537,30 @@ class MeshLlmAgentInjector(BaseInjector):
         is_template = config_dict.get("is_template", False)
         template_path = config_dict.get("template_path")
 
+        # Extract model params (everything except internal config keys)
+        # These are passed through to LiteLLM (e.g., max_tokens, temperature)
+        INTERNAL_CONFIG_KEYS = {
+            "filter",
+            "filter_mode",
+            "provider",
+            "model",
+            "api_key",
+            "max_iterations",
+            "system_prompt",
+            "system_prompt_file",
+            "is_template",
+            "template_path",
+            "context_param",
+            "output_mode",
+        }
+        default_model_params = {
+            k: v for k, v in config_dict.items() if k not in INTERNAL_CONFIG_KEYS
+        }
+        if default_model_params:
+            logger.debug(
+                f"🔧 Extracted default model params for {function_id}: {list(default_model_params.keys())}"
+            )
+
         # Determine vendor for provider handler selection
         # Priority: 1) From registry (mesh delegation), 2) From model name, 3) None
         vendor = llm_agent_data.get("vendor")
@@ -564,6 +588,7 @@ class MeshLlmAgentInjector(BaseInjector):
                 "provider_proxy"
             ),  # Provider proxy for mesh delegation
             vendor=vendor,  # Vendor for provider handler selection (from registry or model name)
+            default_model_params=default_model_params,  # Decorator-level LLM params
         )
 
         logger.debug(
@@ -574,6 +599,11 @@ class MeshLlmAgentInjector(BaseInjector):
                 if isinstance(config_dict.get("provider"), dict)
                 else ""
             )
+            + (
+                f", model_params={list(default_model_params.keys())}"
+                if default_model_params
+                else ""
+            )
         )
 
         return llm_agent

_mcp_mesh/engine/session_aware_client.py

@@ -11,8 +11,8 @@ import logging
 import random
 from typing import Any, Dict, List, Optional
 
-from .mcp_client_proxy import MCPClientProxy
 from .session_manager import get_session_manager
+from .unified_mcp_proxy import UnifiedMCPProxy
 
 logger = logging.getLogger(__name__)
 
@@ -72,10 +72,10 @@ class SessionAwareMCPClient:
         )
 
         # Create client proxy and make the call
-        client = MCPClientProxy(target_agent["endpoint"], capability)
+        proxy = UnifiedMCPProxy(target_agent["endpoint"], capability)
 
         try:
-            result = await client._async_call(**arguments)
+            result = await proxy(**arguments)
 
             # Update session affinity if this was a session-required call
             if session_id and routing_strategy.get("session_required"):

_mcp_mesh/engine/unified_mcp_proxy.py

@@ -1,7 +1,7 @@
 """Unified MCP Proxy using FastMCP's built-in client.
 
-This replaces both MCPClientProxy and FullMCPProxy with a single implementation
-that uses FastMCP's superior client capabilities.
+This is the primary MCP client proxy for cross-service communication,
+using FastMCP's superior client capabilities with async support.
 """
 
 import asyncio
@@ -10,6 +10,13 @@ import uuid
 from collections.abc import AsyncIterator
 from typing import Any, Optional
 
+from ..shared.logging_config import (
+    format_log_value,
+    format_result_summary,
+    get_trace_prefix,
+)
+from ..shared.sse_parser import SSEParser
+
 logger = logging.getLogger(__name__)
 
 
@@ -227,9 +234,6 @@ class UnifiedMCPProxy:
         self.collect_performance_metrics = self.kwargs_config.get(
             "performance_metrics", True
         )
-        self.include_telemetry_in_response = self.kwargs_config.get(
-            "include_telemetry_response", True
-        )
 
         # Agent context collection
         self.collect_agent_context = self.kwargs_config.get(
@@ -397,12 +401,22 @@ class UnifiedMCPProxy:
 
         start_time = time.time()
 
+        # Get trace prefix if available
+        tp = get_trace_prefix()
+
+        # Log cross-agent call - summary line
+        arg_keys = list(arguments.keys()) if arguments else []
+        self.logger.debug(
+            f"{tp}🔄 Cross-agent call: {self.endpoint}/{name} (timeout: {self.timeout}s, args={arg_keys})"
+        )
+        # Log full args (will be TRACE later)
+        self.logger.debug(
+            f"{tp}🔄 Cross-agent call args: {format_log_value(arguments)}"
+        )
+
         try:
             # 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}")
-
-            # Create client with automatic trace header injection
             client_instance = self._create_fastmcp_client(mcp_endpoint)
 
             async with client_instance as client:
@@ -422,21 +436,13 @@ class UnifiedMCPProxy:
                 # Convert CallToolResult to native Python structures for client simplicity
                 converted_result = self._convert_mcp_result_to_python(result)
 
-                # Add telemetry metadata to converted result if enabled
-                if self.include_telemetry_in_response and isinstance(
-                    converted_result, dict
-                ):
-                    converted_result["_telemetry"] = {
-                        "method": "fastmcp_client",
-                        "duration_ms": duration_ms,
-                        "endpoint": mcp_endpoint,
-                        "tool_name": name,
-                        "telemetry_enabled": self.telemetry_enabled,
-                        "distributed_tracing": self.distributed_tracing_enabled,
-                    }
-
+                # Log success - summary line
                 self.logger.info(
-                    f"✅ FastMCP tool call successful: {name} in {duration_ms}ms"
+                    f"{tp}✅ FastMCP tool call successful: {name} in {duration_ms}ms → {format_result_summary(converted_result)}"
+                )
+                # Log full result (will be TRACE later)
+                self.logger.debug(
+                    f"{tp}🔄 Cross-agent call result: {format_log_value(converted_result)}"
                 )
                 return converted_result
 
@@ -556,6 +562,52 @@ class UnifiedMCPProxy:
         else:
             return {"data": str(structured_content)}
 
+    def _normalize_http_result(self, result: Any) -> Any:
+        """Normalize HTTP fallback result to match FastMCP format.
+
+        Extracts structured content from MCP envelope so callers get consistent
+        response format regardless of transport method (FastMCP vs HTTP fallback).
+        """
+        if not isinstance(result, dict):
+            return result
+
+        # If result has "content" array (MCP envelope format), extract the data
+        if "content" in result and isinstance(result["content"], list):
+            content_list = result["content"]
+
+            if not content_list:
+                return None
+
+            # Single content item - extract and parse
+            if len(content_list) == 1:
+                content_item = content_list[0]
+                if isinstance(content_item, dict) and "text" in content_item:
+                    # Try to parse JSON text content
+                    try:
+                        import json
+
+                        return json.loads(content_item["text"])
+                    except (json.JSONDecodeError, TypeError):
+                        return content_item["text"]
+                return content_item
+
+            # Multiple content items - convert each
+            converted_items = []
+            for item in content_list:
+                if isinstance(item, dict) and "text" in item:
+                    try:
+                        import json
+
+                        converted_items.append(json.loads(item["text"]))
+                    except (json.JSONDecodeError, TypeError):
+                        converted_items.append(item["text"])
+                else:
+                    converted_items.append(item)
+            return {"content": converted_items, "type": "multi_content"}
+
+        # Result is already in clean format
+        return result
+
     async def _fallback_http_call(self, name: str, arguments: dict = None) -> Any:
         """Enhanced fallback HTTP call using httpx directly with performance tracking."""
         import time
@@ -614,41 +666,10 @@ class UnifiedMCPProxy:
                     f"📄 Response length: {len(response_text)} chars, starts with: {response_text[:100]}"
                 )
 
-                data = None
-
-                # Handle SSE format
-                if response_text.startswith("event:") or "data:" in response_text:
-                    self.logger.debug("🔄 Parsing SSE format response")
-                    # Parse SSE format - handle multiple events
-                    for line in response_text.split("\n"):
-                        line = line.strip()
-                        if line.startswith("data:"):
-                            json_str = line[5:].strip()
-                            if json_str and json_str != "":
-                                try:
-                                    data = json.loads(json_str)
-                                    self.logger.debug(
-                                        f"✅ Successfully parsed SSE data: {type(data)}"
-                                    )
-                                    break
-                                except json.JSONDecodeError as e:
-                                    self.logger.warning(
-                                        f"⚠️ Failed to parse SSE line: {json_str[:100]}, error: {e}"
-                                    )
-                                    continue
-                else:
-                    # Plain JSON response
-                    self.logger.debug("🔄 Parsing plain JSON response")
-                    try:
-                        data = json.loads(response_text)
-                    except json.JSONDecodeError as e:
-                        self.logger.error(
-                            f"❌ Failed to parse JSON response: {e}, content: {response_text[:200]}"
-                        )
-                        raise RuntimeError(f"Invalid JSON response: {e}")
-
-                if data is None:
-                    raise RuntimeError("No valid data found in response")
+                # Use shared SSE parser for both SSE and plain JSON responses
+                data = SSEParser.parse_sse_response(
+                    response_text, f"UnifiedMCPProxy.{name}"
+                )
 
                 # Check for JSON-RPC error
                 if "error" in data:
@@ -673,39 +694,10 @@ class UnifiedMCPProxy:
                     f"📊 HTTP fallback performance: {duration_ms}ms for tool '{name}'"
                 )
 
-                if isinstance(result, dict) and "content" in result:
-                    self.logger.info(
-                        f"✅ Successfully parsed result with content in {duration_ms}ms"
-                    )
-                    # Add performance metadata to result if enabled
-                    if self.include_telemetry_in_response and isinstance(result, dict):
-                        result["_telemetry"] = {
-                            "method": "http_fallback",
-                            "duration_ms": duration_ms,
-                            "endpoint": self.endpoint,
-                            "tool_name": name,
-                            "telemetry_enabled": self.telemetry_enabled,
-                            "fallback_reason": "fastmcp_unavailable",
-                        }
-                    return result
-                else:
-                    # Wrap in CallToolResult-like structure with performance data
-                    self.logger.info(
-                        f"🔄 Wrapping result in CallToolResult structure ({duration_ms}ms)"
-                    )
-                    wrapped_result = {
-                        "content": [{"type": "text", "text": str(result)}]
-                    }
-                    if self.include_telemetry_in_response:
-                        wrapped_result["_telemetry"] = {
-                            "method": "http_fallback",
-                            "duration_ms": duration_ms,
-                            "endpoint": self.endpoint,
-                            "tool_name": name,
-                            "telemetry_enabled": self.telemetry_enabled,
-                            "fallback_reason": "fastmcp_unavailable",
-                        }
-                    return wrapped_result
+                # Normalize HTTP response to match FastMCP format
+                normalized_result = self._normalize_http_result(result)
+                self.logger.info(f"✅ HTTP fallback successful in {duration_ms}ms")
+                return normalized_result
 
         except ImportError:
             raise RuntimeError("httpx not available for HTTP fallback")

_mcp_mesh/pipeline/api_heartbeat/api_dependency_resolution.py

@@ -215,11 +215,6 @@ class APIDependencyResolutionStep(PipelineStep):
 
             # Import here to avoid circular imports
             from ...engine.dependency_injector import get_global_injector
-            from ...engine.full_mcp_proxy import EnhancedFullMCPProxy, FullMCPProxy
-            from ...engine.mcp_client_proxy import (
-                EnhancedMCPClientProxy,
-                MCPClientProxy,
-            )
 
             injector = get_global_injector()
 
@@ -421,87 +416,3 @@ class APIDependencyResolutionStep(PipelineStep):
                 f"❌ Failed to process API heartbeat response for rewiring: {e}"
             )
             # Don't raise - this should not break the heartbeat loop
-
-    def _determine_api_proxy_type_for_capability(
-        self, capability: str, injector
-    ) -> str:
-        """
-        Determine which proxy type to use for API route handlers.
-
-        Since McpAgent has been removed, all API route handlers now use MCPClientProxy
-        for McpMeshAgent parameters.
-
-        Args:
-            capability: The capability name to check
-            injector: The dependency injector instance
-
-        Returns:
-            "MCPClientProxy"
-        """
-        # Note: This method always returns "MCPClientProxy" since McpAgent was removed.
-        # All McpMeshAgent parameters use MCPClientProxy.
-        self.logger.debug(
-            f"🔍 API route handlers for capability '{capability}' → using MCPClientProxy"
-        )
-        return "MCPClientProxy"
-
-    def _create_proxy_for_api(
-        self,
-        proxy_type: str,
-        endpoint: str,
-        dep_function_name: str,
-        kwargs_config: dict,
-    ):
-        """
-        Create the appropriate proxy instance for API route handlers.
-
-        Args:
-            proxy_type: "FullMCPProxy" or "MCPClientProxy"
-            endpoint: Target endpoint URL
-            dep_function_name: Target function name
-            kwargs_config: Additional configuration (timeout, retry, etc.)
-
-        Returns:
-            Proxy instance
-        """
-        from ...engine.full_mcp_proxy import EnhancedFullMCPProxy, FullMCPProxy
-        from ...engine.mcp_client_proxy import EnhancedMCPClientProxy, MCPClientProxy
-
-        if proxy_type == "FullMCPProxy":
-            # Use enhanced proxy if kwargs available
-            if kwargs_config:
-                proxy = EnhancedFullMCPProxy(
-                    endpoint,
-                    dep_function_name,
-                    kwargs_config=kwargs_config,
-                )
-                self.logger.debug(
-                    f"🔧 Created EnhancedFullMCPProxy for API with kwargs: {kwargs_config}"
-                )
-            else:
-                proxy = FullMCPProxy(
-                    endpoint,
-                    dep_function_name,
-                    kwargs_config=kwargs_config,
-                )
-                self.logger.debug("🔧 Created FullMCPProxy for API (no kwargs)")
-            return proxy
-        else:
-            # Use enhanced proxy if kwargs available
-            if kwargs_config:
-                proxy = EnhancedMCPClientProxy(
-                    endpoint,
-                    dep_function_name,
-                    kwargs_config=kwargs_config,
-                )
-                self.logger.debug(
-                    f"🔧 Created EnhancedMCPClientProxy for API with kwargs: {kwargs_config}"
-                )
-            else:
-                proxy = MCPClientProxy(
-                    endpoint,
-                    dep_function_name,
-                    kwargs_config=kwargs_config,
-                )
-                self.logger.debug("🔧 Created MCPClientProxy for API (no kwargs)")
-            return proxy

_mcp_mesh/pipeline/api_heartbeat/api_fast_heartbeat_check.py

@@ -42,7 +42,7 @@ class APIFastHeartbeatStep(PipelineStep):
         Returns:
             PipelineResult with fast_heartbeat_status in context
         """
-        self.logger.debug("Starting API fast heartbeat check...")
+        self.logger.trace("Starting API fast heartbeat check...")
 
         result = PipelineResult(message="API fast heartbeat check completed")
 
@@ -57,7 +57,7 @@ class APIFastHeartbeatStep(PipelineStep):
             if not registry_wrapper:
                 raise ValueError("registry_wrapper is required in context")
 
-            self.logger.debug(
+            self.logger.trace(
                 f"🚀 Performing API fast heartbeat check for service '{service_id}'"
             )
 
@@ -114,4 +114,4 @@ class APIFastHeartbeatStep(PipelineStep):
             if key not in result.context:
                 result.add_context(key, value)
 
-        return result
+        return result