solace-agent-mesh 1.7.1__py3-none-any.whl → 1.13.2__py3-none-any.whl

solace_agent_mesh/agent/adk/embed_resolving_mcp_toolset.py

@@ -10,6 +10,7 @@ The base class is determined at import time based on enterprise package availabi
 """
 
 import logging
+import time
 from typing import Any
 
 from google.adk.auth.credential_manager import CredentialManager
@@ -28,6 +29,7 @@ from ..utils.context_helpers import get_original_session_id
 
 log = logging.getLogger(__name__)
 
+
 def _get_base_mcp_toolset_class() -> tuple[type[MCPToolset], bool]:
     """
     Factory function to determine which base MCP toolset class to use for inheritance.
@@ -76,6 +78,62 @@ def _get_base_mcp_tool_class() -> tuple[type[MCPTool], bool]:
 _BaseMcpToolClass, _base_supports_tool_config = _get_base_mcp_tool_class()
 
 
+def _log_mcp_tool_call(userId, agentId, tool_name, session_id):
+    """A short log message so that customers can track tool usage per user/agent"""
+    log.info(
+        "MCP Tool Call - UserID: %s, AgentID: %s, ToolName: %s, SessionID: %s",
+        userId,
+        agentId,
+        tool_name,
+        session_id,
+        extra={
+            "user_id": userId,
+            "agent_id": agentId,
+            "tool_name": tool_name,
+            "session_id": session_id,
+        },
+    )
+
+
+def _log_mcp_tool_success(userId, agentId, tool_name, session_id, duration_ms):
+    """A short log message so that customers can track successful tool completion per user/agent"""
+    log.info(
+        "MCP Tool Success - UserID: %s, AgentID: %s, ToolName: %s, SessionID: %s, Duration: %.2fms",
+        userId,
+        agentId,
+        tool_name,
+        session_id,
+        duration_ms,
+        extra={
+            "user_id": userId,
+            "agent_id": agentId,
+            "tool_name": tool_name,
+            "session_id": session_id,
+            "duration_ms": duration_ms,
+        },
+    )
+
+
+def _log_mcp_tool_failure(userId, agentId, tool_name, session_id, duration_ms, error):
+    """A short log message so that customers can track tool failures per user/agent"""
+    log.error(
+        "MCP Tool Failure - UserID: %s, AgentID: %s, ToolName: %s, SessionID: %s, Duration: %.2fms, Error: %s",
+        userId,
+        agentId,
+        tool_name,
+        session_id,
+        duration_ms,
+        str(error),
+        extra={
+            "user_id": userId,
+            "agent_id": agentId,
+            "tool_name": tool_name,
+            "session_id": session_id,
+            "duration_ms": duration_ms,
+        },
+    )
+
+
 class EmbedResolvingMCPTool(_BaseMcpToolClass):
     """
     Custom MCPTool that resolves embeds in parameters before calling the actual MCP tool.
@@ -112,7 +170,7 @@ class EmbedResolvingMCPTool(_BaseMcpToolClass):
                     original_mcp_tool._mcp_tool, "auth_credential", None
                 ),
             )
-        self._original_mcp_tool = original_mcp_tool
+        self._original_mcp_tool: MCPTool = original_mcp_tool
         self._tool_config = tool_config or {}
 
     async def _resolve_embeds_recursively(
@@ -275,6 +333,38 @@ class EmbedResolvingMCPTool(_BaseMcpToolClass):
         )
         return data
 
+
+    async def _execute_tool_with_audit_logs(self, tool_call, tool_context):
+        _log_mcp_tool_call(
+            tool_context.session.user_id,
+            tool_context.agent_name,
+            self.name,
+            tool_context.session.id,
+        )
+        start_time = time.perf_counter()
+        try:
+            result = await tool_call()
+            duration_ms = (time.perf_counter() - start_time) * 1000
+            _log_mcp_tool_success(
+                tool_context.session.user_id,
+                tool_context.agent_name,
+                self.name,
+                tool_context.session.id,
+                duration_ms,
+            )
+            return result
+        except Exception as e:
+            duration_ms = (time.perf_counter() - start_time) * 1000
+            _log_mcp_tool_failure(
+                tool_context.session.user_id,
+                tool_context.agent_name,
+                self.name,
+                tool_context.session.id,
+                duration_ms,
+                e,
+            )
+            raise
+
     async def _run_async_impl(
         self, *, args, tool_context: ToolContext, credential
     ) -> Any:
@@ -285,7 +375,6 @@ class EmbedResolvingMCPTool(_BaseMcpToolClass):
 
         # Get context for embed resolution - pass the tool_context object directly
         context_for_embeds = tool_context
-
         if context_for_embeds:
             log.debug(
                 "%s Starting recursive embed resolution for all parameters. Context type: %s",
@@ -321,13 +410,14 @@ class EmbedResolvingMCPTool(_BaseMcpToolClass):
                 log_identifier,
             )
             resolved_args = args
-
         # Call the original MCP tool with resolved parameters
-        return await self._original_mcp_tool._run_async_impl(
-            args=resolved_args, tool_context=tool_context, credential=credential
+        return await self._execute_tool_with_audit_logs(
+            lambda: self._original_mcp_tool._run_async_impl(
+                args=resolved_args, tool_context=tool_context, credential=credential
+            ),
+            tool_context,
         )
 
-
 # Get the base toolset class to use for inheritance
 _BaseMcpToolsetClass, _base_toolset_supports_tool_config = _get_base_mcp_toolset_class()
 
@@ -343,6 +433,7 @@ class EmbedResolvingMCPToolset(_BaseMcpToolsetClass):
         self,
         connection_params,
         tool_filter=None,
+        tool_name_prefix=None,
         auth_scheme=None,
         auth_credential=None,
         auth_discovery=None,
@@ -357,6 +448,7 @@ class EmbedResolvingMCPToolset(_BaseMcpToolsetClass):
             super().__init__(
                 connection_params=connection_params,
                 tool_filter=tool_filter,
+                tool_name_prefix=tool_name_prefix,
                 auth_scheme=auth_scheme,
                 auth_credential=auth_credential,
                 auth_discovery=auth_discovery,
@@ -367,6 +459,7 @@ class EmbedResolvingMCPToolset(_BaseMcpToolsetClass):
             super().__init__(
                 connection_params=connection_params,
                 tool_filter=tool_filter,
+                tool_name_prefix=tool_name_prefix,
                 auth_scheme=auth_scheme,
                 auth_credential=auth_credential,
             )
@@ -403,4 +496,3 @@ class EmbedResolvingMCPToolset(_BaseMcpToolsetClass):
 
         self._tool_cache = embed_resolving_tools
         return embed_resolving_tools
-

solace_agent_mesh/agent/adk/intelligent_mcp_callbacks.py

@@ -19,6 +19,7 @@ from .mcp_content_processor import MCPContentProcessor, MCPContentProcessorConfi
 from ...agent.utils.artifact_helpers import (
     save_artifact_with_metadata,
     DEFAULT_SCHEMA_MAX_KEYS,
+    DEFAULT_SCHEMA_INFERENCE_DEPTH,
 )
 from ...agent.utils.context_helpers import get_original_session_id
 
@@ -28,6 +29,41 @@ if TYPE_CHECKING:
     from ...agent.sac.component import SamAgentComponent
 
 
+def _get_schema_config_from_tool_or_agent(
+    tool: BaseTool,
+    host_component: "SamAgentComponent",
+    config_key: str,
+    default_value: Any,
+) -> Any:
+    """
+    Get schema configuration from tool_config if available, otherwise fall back to agent config.
+
+    This allows per-tool override of schema settings like schema_inference_depth.
+
+    Args:
+        tool: The MCP tool instance (may have _tool_config attribute)
+        host_component: The agent component for fallback config
+        config_key: The configuration key to look up
+        default_value: Default value if not found in either location
+
+    Returns:
+        The configuration value from tool config, agent config, or default
+    """
+    # Check if tool has tool_config with this setting
+    tool_config = getattr(tool, "_tool_config", None)
+    if tool_config and isinstance(tool_config, dict):
+        if config_key in tool_config:
+            log.debug(
+                "Using per-tool config for %s: %s",
+                config_key,
+                tool_config[config_key],
+            )
+            return tool_config[config_key]
+
+    # Fall back to agent-level config
+    return host_component.get_config(config_key, default_value)
+
+
 class McpSaveStatus(str, Enum):
     """Enumeration for the status of an MCP save operation."""
 
@@ -164,7 +200,7 @@ async def save_mcp_response_as_artifact_intelligent(
                 if hasattr(item, "uri"):
                     item.uri = str(item.uri)
                 result_dict = await _save_content_item_as_artifact(
-                    item, tool_context, host_component
+                    item, tool, tool_context, host_component
                 )
                 if result_dict.get("status") in ["success", "partial_success"]:
                     saved_artifacts.append(SavedArtifactInfo(**result_dict))
@@ -279,6 +315,7 @@ async def save_mcp_response_as_artifact_intelligent(
 
 async def _save_content_item_as_artifact(
     content_item,
+    tool: BaseTool,
     tool_context: ToolContext,
     host_component: "SamAgentComponent",
 ) -> Dict[str, Any]:
@@ -294,8 +331,12 @@ async def _save_content_item_as_artifact(
         app_name = host_component.agent_name
         user_id = tool_context._invocation_context.user_id
         session_id = get_original_session_id(tool_context._invocation_context)
-        schema_max_keys = host_component.get_config(
-            "schema_max_keys", DEFAULT_SCHEMA_MAX_KEYS
+        # Get schema config from tool_config (per-tool) or agent config (fallback)
+        schema_max_keys = _get_schema_config_from_tool_or_agent(
+            tool, host_component, "schema_max_keys", DEFAULT_SCHEMA_MAX_KEYS
+        )
+        schema_inference_depth = _get_schema_config_from_tool_or_agent(
+            tool, host_component, "schema_inference_depth", DEFAULT_SCHEMA_INFERENCE_DEPTH
         )
         artifact_timestamp = datetime.now(timezone.utc)
 
@@ -318,6 +359,7 @@ async def _save_content_item_as_artifact(
             metadata_dict=content_item.metadata,
             timestamp=artifact_timestamp,
             schema_max_keys=schema_max_keys,
+            schema_inference_depth=schema_inference_depth,
             tool_context=tool_context,
         )
 
@@ -376,8 +418,12 @@ async def _save_raw_mcp_response_fallback(
         app_name = host_component.agent_name
         user_id = tool_context._invocation_context.user_id
         session_id = get_original_session_id(tool_context._invocation_context)
-        schema_max_keys = host_component.get_config(
-            "schema_max_keys", DEFAULT_SCHEMA_MAX_KEYS
+        # Get schema config from tool_config (per-tool) or agent config (fallback)
+        schema_max_keys = _get_schema_config_from_tool_or_agent(
+            tool, host_component, "schema_max_keys", DEFAULT_SCHEMA_MAX_KEYS
+        )
+        schema_inference_depth = _get_schema_config_from_tool_or_agent(
+            tool, host_component, "schema_inference_depth", DEFAULT_SCHEMA_INFERENCE_DEPTH
         )
 
         save_result = await save_artifact_with_metadata(
@@ -391,6 +437,7 @@ async def _save_raw_mcp_response_fallback(
             metadata_dict=metadata_for_saving,
             timestamp=artifact_timestamp,
             schema_max_keys=schema_max_keys,
+            schema_inference_depth=schema_inference_depth,
             tool_context=tool_context,
         )
 

solace_agent_mesh/agent/adk/mcp_content_processor.py

@@ -540,7 +540,7 @@ class MCPContentProcessor:
 
         format_mapping = {
             TextFormat.JSON: ("application/json", ".json"),
-            TextFormat.YAML: ("application/x-yaml", ".yaml"),
+            TextFormat.YAML: ("application/yaml", ".yaml"),
             TextFormat.CSV: ("text/csv", ".csv"),
             TextFormat.MARKDOWN: ("text/markdown", ".md"),
             TextFormat.PLAIN: ("text/plain", ".txt"),

solace_agent_mesh/agent/adk/models/lite_llm.py

@@ -328,26 +328,41 @@ def _schema_to_dict(schema: types.Schema) -> dict:
     """
 
     schema_dict = schema.model_dump(exclude_none=True)
+
+    # Convert top-level type from enum to lowercase string
     if "type" in schema_dict:
-        schema_dict["type"] = schema_dict["type"].lower()
+        if isinstance(schema_dict["type"], types.Type):
+            schema_dict["type"] = schema_dict["type"].value.lower()
+        else:
+            schema_dict["type"] = str(schema_dict["type"]).lower()
+
+    # Recursively handle items (for array types)
     if "items" in schema_dict:
-        if isinstance(schema_dict["items"], dict):
+        # Check if we have the original Schema object for items
+        if isinstance(schema.items, types.Schema):
+            # Recursively convert the Schema object - this ensures nested Type enums are converted
+            schema_dict["items"] = _schema_to_dict(schema.items)
+        elif isinstance(schema_dict["items"], dict):
+            # If items is already a dict, validate and recurse
             schema_dict["items"] = _schema_to_dict(
                 types.Schema.model_validate(schema_dict["items"])
             )
-        elif isinstance(schema_dict["items"]["type"], types.Type):
-            schema_dict["items"]["type"] = TYPE_LABELS[
-                schema_dict["items"]["type"].value
-            ]
+
+    # Recursively handle properties (for object types)
     if "properties" in schema_dict:
         properties = {}
         for key, value in schema_dict["properties"].items():
             if isinstance(value, types.Schema):
+                # If it's a Schema object, recursively convert it
                 properties[key] = _schema_to_dict(value)
+            elif isinstance(value, dict):
+                # If it's already a dict, validate and recurse to handle nested Type enums
+                properties[key] = _schema_to_dict(
+                    types.Schema.model_validate(value)
+                )
             else:
+                # For other types, just copy as-is
                 properties[key] = value
-                if "type" in properties[key]:
-                    properties[key]["type"] = properties[key]["type"].lower()
         schema_dict["properties"] = properties
     return schema_dict
 
@@ -366,20 +381,23 @@ def _function_declaration_to_tool_param(
 
     assert function_declaration.name
 
-    properties = {}
-    if function_declaration.parameters and function_declaration.parameters.properties:
-        for key, value in function_declaration.parameters.properties.items():
-            properties[key] = _schema_to_dict(value)
+    # Convert the entire parameters schema to ensure all fields (type, properties, required, etc.)
+    # are properly converted, including nested Type enums
+    # If no parameters provided, default to empty object schema (required by OpenAI)
+    if function_declaration.parameters:
+        parameters = _schema_to_dict(function_declaration.parameters)
+    else:
+        parameters = {
+            "type": "object",
+            "properties": {},
+        }
 
     return {
         "type": "function",
         "function": {
             "name": function_declaration.name,
             "description": function_declaration.description or "",
-            "parameters": {
-                "type": "object",
-                "properties": properties,
-            },
+            "parameters": parameters,
         },
     }
 

solace_agent_mesh/agent/adk/models/oauth2_token_manager.py

@@ -6,12 +6,9 @@ It handles token acquisition, caching, and automatic refresh with proper error h
 
 import asyncio
 import logging
-import random
-import time
-from typing import Any, Dict, Optional
-
-import httpx
+from typing import Optional
 
+from solace_agent_mesh.common.oauth import OAuth2RetryClient, is_token_expired
 from solace_agent_mesh.common.utils.in_memory_cache import InMemoryCache
 
 logger = logging.getLogger(__name__)
@@ -72,17 +69,19 @@ class OAuth2ClientCredentialsTokenManager:
         self.scope = scope
         self.ca_cert_path = ca_cert_path
         self.refresh_buffer_seconds = refresh_buffer_seconds
-        self.max_retries = max_retries
-        
+
         # Thread-safe token access
         self._lock = asyncio.Lock()
-        
+
         # Token cache using existing InMemoryCache singleton
         self._cache = InMemoryCache()
-        
+
         # Cache key for this token manager instance
         self._cache_key = f"oauth_token_{hash((token_url, client_id))}"
-        
+
+        # OAuth client with retry logic
+        self._oauth_client = OAuth2RetryClient(max_retries=max_retries)
+
         logger.info(
             "OAuth2ClientCredentialsTokenManager initialized for endpoint: %s",
             token_url
@@ -104,14 +103,24 @@ class OAuth2ClientCredentialsTokenManager:
         async with self._lock:
             # Check if we have a cached token
             cached_token_data = self._cache.get(self._cache_key)
-            
-            if cached_token_data and not self._is_token_expired(cached_token_data):
+
+            if cached_token_data and not is_token_expired(
+                cached_token_data["expires_at"], buffer_seconds=self.refresh_buffer_seconds
+            ):
                 logger.debug("Using cached OAuth token")
                 return cached_token_data["access_token"]
-            
-            # Fetch new token
+
+            # Fetch new token using common OAuth client
             logger.info("Fetching new OAuth token from %s", self.token_url)
-            token_data = await self._fetch_token()
+            verify = self.ca_cert_path if self.ca_cert_path else True
+            token_data = await self._oauth_client.fetch_client_credentials_token(
+                token_url=self.token_url,
+                client_id=self.client_id,
+                client_secret=self.client_secret,
+                scope=self.scope,
+                verify=verify,
+                timeout=30.0,
+            )
             
             # Cache the token with TTL
             expires_in = token_data.get("expires_in", 3600)  # Default 1 hour
@@ -121,125 +130,3 @@ class OAuth2ClientCredentialsTokenManager:
             
             logger.info("OAuth token cached with TTL: %d seconds", cache_ttl)
             return token_data["access_token"]
-
-    def _is_token_expired(self, token_data: Dict[str, Any]) -> bool:
-        """Check if a token is expired or near expiry.
-        
-        Args:
-            token_data: Token data dictionary with 'expires_at' timestamp
-            
-        Returns:
-            True if token is expired or near expiry, False otherwise
-        """
-        if "expires_at" not in token_data:
-            return True
-            
-        current_time = time.time()
-        expires_at = token_data["expires_at"]
-        
-        # Consider token expired if it expires within the buffer time
-        return current_time >= (expires_at - self.refresh_buffer_seconds)
-
-    async def _fetch_token(self) -> Dict[str, Any]:
-        """Fetch a new OAuth 2.0 access token from the token endpoint.
-
-        Implements retry logic with exponential backoff for transient failures.
-
-        Returns:
-            Token data dictionary containing access_token, expires_in, etc.
-
-        Raises:
-            httpx.HTTPError: If HTTP request fails after all retries
-            ValueError: If response is invalid or missing required fields
-        """
-        # Prepare request payload
-        payload = {
-            "grant_type": "client_credentials",
-            "client_id": self.client_id,
-            "client_secret": self.client_secret,
-        }
-        
-        if self.scope:
-            payload["scope"] = self.scope
-        
-        # Configure HTTP client with SSL settings
-        verify = True
-        if self.ca_cert_path:
-            verify = self.ca_cert_path
-            
-        headers = {
-            "Content-Type": "application/x-www-form-urlencoded",
-            "Accept": "application/json",
-        }
-        
-        last_exception = None
-
-        for attempt in range(self.max_retries + 1):
-            try:
-                async with httpx.AsyncClient(verify=verify) as client:
-                    response = await client.post(
-                        self.token_url,
-                        data=payload,
-                        headers=headers,
-                        timeout=30.0,
-                    )
-                    response.raise_for_status()
-
-                    token_data = response.json()
-
-                    # Validate response
-                    if "access_token" not in token_data:
-                        raise ValueError("Token response missing 'access_token' field")
-
-                    # Add expiration timestamp for cache management
-                    expires_in = token_data.get("expires_in", 3600)
-                    token_data["expires_at"] = time.time() + expires_in
-
-                    logger.info("Successfully fetched OAuth token, expires in %d seconds", expires_in)
-                    return token_data
-
-            except httpx.HTTPStatusError as e:
-                last_exception = e
-                # Don't retry on 4xx errors (client errors)
-                if 400 <= e.response.status_code < 500:
-                    logger.error(
-                        "OAuth token request failed with client error %d: %s",
-                        e.response.status_code,
-                        e.response.text
-                    )
-                    raise
-
-                logger.warning(
-                    "OAuth token request failed with status %d (attempt %d/%d): %s",
-                    e.response.status_code,
-                    attempt + 1,
-                    self.max_retries + 1,
-                    e.response.text
-                )
-
-            except httpx.RequestError as e:
-                last_exception = e
-                logger.warning(
-                    "OAuth token request failed (attempt %d/%d): %s",
-                    attempt + 1,
-                    self.max_retries + 1,
-                    str(e)
-                )
-
-            except Exception as e:
-                last_exception = e
-                logger.error("Unexpected error during OAuth token fetch: %s", str(e))
-                raise
-
-            # Exponential backoff with jitter for retries
-            if attempt < self.max_retries:
-                delay = (2 ** attempt) + random.uniform(0, 1)
-                logger.info("Retrying OAuth token request in %.2f seconds", delay)
-                await asyncio.sleep(delay)
-
-        # All retries exhausted
-        logger.error("OAuth token request failed after %d attempts", self.max_retries + 1)
-        if last_exception:
-            raise last_exception
-        else:
-            raise RuntimeError("OAuth token request failed after all retries")