solace-agent-mesh 1.6.1__py3-none-any.whl → 1.6.2__py3-none-any.whl

solace_agent_mesh/agent/adk/app_llm_agent.py

@@ -3,7 +3,10 @@ Custom LlmAgent subclass for the A2A Host Component.
 """
 
 from typing import Any
+
 from google.adk.agents import LlmAgent
+from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow
+from google.adk.flows.llm_flows.single_flow import SingleFlow
 from pydantic import Field
 
 
@@ -24,3 +27,26 @@ class AppLlmAgent(LlmAgent):
     This field is excluded from Pydantic's model serialization and validation
     if not provided during instantiation, and is intended to be set post-init.
     """
+
+    # Override the _llm_flow property to inject the patched auth preprocessor
+    @property
+    def _llm_flow(self) -> BaseLlmFlow:
+        try:
+            from solace_agent_mesh_enterprise.auth.auth_preprocessor_patch import (
+                request_processor,
+            )
+        except ImportError:
+            # If enterprise module doesn't exist, use standard parent flow
+            return super()._llm_flow
+
+        llm_flow = super()._llm_flow
+        if isinstance(llm_flow, SingleFlow):
+            # Replace auth_preprocessor.request_processor with the patched version from _AuthLlmRequestProcessorPatched
+            from google.adk.auth import auth_preprocessor
+
+            for i, processor in enumerate(llm_flow.request_processors):
+                if processor is auth_preprocessor.request_processor:
+                    llm_flow.request_processors[i] = request_processor
+                    break
+
+        return llm_flow

solace_agent_mesh/agent/adk/artifacts/filesystem_artifact_service.py

@@ -45,7 +45,7 @@ class FilesystemArtifactService(BaseArtifactService):
         try:
             os.makedirs(self.base_path, exist_ok=True)
             logger.info(
-                "FilesystemArtifactService initialized. Base path: %s",
+                "Initialized FilesystemArtifactService. Base path: %s",
                 self.base_path,
             )
         except OSError as e:

solace_agent_mesh/agent/adk/embed_resolving_mcp_toolset.py

@@ -1,44 +1,116 @@
 """
 Custom MCPToolset that resolves embeds in tool parameters before calling MCP tools.
+
+This module uses dynamic inheritance to support both standard and enterprise MCP tools:
+- Standard mode: Inherits from MCPTool and MCPToolset (Google ADK)
+- Enterprise mode: Inherits from McpToolWithManifest and McpToolsetWithManifest
+  (adds manifest support and tool_config parameter)
+
+The base class is determined at import time based on enterprise package availability.
 """
 
 import logging
-import asyncio
-from typing import Dict, List, Optional, Any
-
-from google.adk.tools.mcp_tool import MCPToolset, MCPTool
-from google.adk.tools.mcp_tool.mcp_session_manager import (
-    SseConnectionParams,
-    StdioConnectionParams,
-    StreamableHTTPConnectionParams,
-)
-from google.adk.tools.tool_context import ToolContext
+from typing import Any
 
+from google.adk.auth.credential_manager import CredentialManager
+from google.adk.tools.mcp_tool import MCPTool, MCPToolset
+from google.adk.tools.tool_context import ToolContext
 
-from ..utils.context_helpers import get_original_session_id
 from ...common.utils.embeds import (
-    resolve_embeds_in_string,
-    evaluate_embed,
     EARLY_EMBED_TYPES,
-    LATE_EMBED_TYPES,
     EMBED_DELIMITER_OPEN,
+    LATE_EMBED_TYPES,
+    evaluate_embed,
+    resolve_embeds_in_string,
 )
+from ..utils.context_helpers import get_original_session_id
 
 log = logging.getLogger(__name__)
 
-class EmbedResolvingMCPTool(MCPTool):
+def _get_base_mcp_toolset_class() -> tuple[type[MCPToolset], bool]:
+    """
+    Factory function to determine which base MCP toolset class to use for inheritance.
+
+    Tries to import McpToolsetWithManifest from solace_agent_mesh_enterprise.common
+    and returns it if available. Falls back to base MCPToolset if not available.
+
+    Returns:
+        Tuple of (class, supports_tool_config_flag) where:
+        - class: The base MCPToolset class to inherit from
+        - supports_tool_config_flag: Whether the class supports tool_config parameter
+    """
+    try:
+        from solace_agent_mesh_enterprise.auth.mcp_toolset_with_manifest import (
+            McpToolsetWithManifest,
+        )
+
+        return (McpToolsetWithManifest, True)
+    except ImportError:
+        return (MCPToolset, False)
+
+
+def _get_base_mcp_tool_class() -> tuple[type[MCPTool], bool]:
+    """
+    Factory function to determine which base MCP tool class to use for inheritance.
+
+    Tries to import McpToolWithManifest from solace_agent_mesh_enterprise.common
+    and returns it if available. Falls back to base MCPTool if not available.
+
+    Returns:
+        Tuple of (class, supports_tool_config_flag) where:
+        - class: The base MCPTool class to inherit from
+        - supports_tool_config_flag: Whether the class supports tool_config parameter
+    """
+    try:
+        from solace_agent_mesh_enterprise.auth.mcp_toolset_with_manifest import (
+            McpToolWithManifest,
+        )
+
+        return (McpToolWithManifest, True)
+    except ImportError:
+        return (MCPTool, False)
+
+
+# Get the base tool class to use for inheritance
+_BaseMcpToolClass, _base_supports_tool_config = _get_base_mcp_tool_class()
+
+
+class EmbedResolvingMCPTool(_BaseMcpToolClass):
     """
     Custom MCPTool that resolves embeds in parameters before calling the actual MCP tool.
+    Uses dynamic inheritance to conditionally inherit from McpToolWithManifest when available,
+    falling back to the standard MCPTool base class.
     """
 
-    def __init__(self, original_mcp_tool: MCPTool, tool_config: Optional[Dict] = None):
+    def __init__(
+        self,
+        original_mcp_tool: MCPTool,
+        tool_config: dict | None = None,
+        credential_manager: CredentialManager | None = None,
+    ):
         # Copy all attributes from the original tool
-        super().__init__(
-            mcp_tool=original_mcp_tool._mcp_tool,
-            mcp_session_manager=original_mcp_tool._mcp_session_manager,
-            auth_scheme=getattr(original_mcp_tool, "_auth_scheme", None),
-            auth_credential=getattr(original_mcp_tool, "_auth_credential", None),
-        )
+        if _base_supports_tool_config:
+            super().__init__(
+                mcp_tool=original_mcp_tool._mcp_tool,
+                mcp_session_manager=original_mcp_tool._mcp_session_manager,
+                auth_scheme=getattr(original_mcp_tool._mcp_tool, "auth_scheme", None),
+                auth_credential=getattr(
+                    original_mcp_tool._mcp_tool, "auth_credential", None
+                ),
+                auth_discovery=getattr(
+                    original_mcp_tool._mcp_tool, "auth_discovery", None
+                ),
+                credential_manager=credential_manager,
+            )
+        else:
+            super().__init__(
+                mcp_tool=original_mcp_tool._mcp_tool,
+                mcp_session_manager=original_mcp_tool._mcp_session_manager,
+                auth_scheme=getattr(original_mcp_tool._mcp_tool, "auth_scheme", None),
+                auth_credential=getattr(
+                    original_mcp_tool._mcp_tool, "auth_credential", None
+                ),
+            )
         self._original_mcp_tool = original_mcp_tool
         self._tool_config = tool_config or {}
 
@@ -254,9 +326,15 @@ class EmbedResolvingMCPTool(MCPTool):
         )
 
 
-class EmbedResolvingMCPToolset(MCPToolset):
+# Get the base toolset class to use for inheritance
+_BaseMcpToolsetClass, _base_toolset_supports_tool_config = _get_base_mcp_toolset_class()
+
+
+class EmbedResolvingMCPToolset(_BaseMcpToolsetClass):
     """
     Custom MCPToolset that creates EmbedResolvingMCPTool instances for embed resolution.
+    Uses dynamic inheritance to conditionally inherit from McpToolsetWithManifest when available,
+    falling back to the standard MCPToolset base class.
     """
 
     def __init__(
@@ -265,20 +343,43 @@ class EmbedResolvingMCPToolset(MCPToolset):
         tool_filter=None,
         auth_scheme=None,
         auth_credential=None,
-        tool_config: Optional[Dict] = None,
+        auth_discovery=None,
+        tool_config: dict | None = None,
+        credential_manager: CredentialManager | None = None,
     ):
-        super().__init__(
-            connection_params=connection_params,
-            tool_filter=tool_filter,
-            auth_scheme=auth_scheme,
-            auth_credential=auth_credential,
-        )
+        # Store tool_config for later use
         self._tool_config = tool_config or {}
 
-    async def get_tools(self, readonly_context=None) -> List[MCPTool]:
+        # Initialize base class with appropriate parameters
+        if _base_toolset_supports_tool_config:
+            super().__init__(
+                connection_params=connection_params,
+                tool_filter=tool_filter,
+                auth_scheme=auth_scheme,
+                auth_credential=auth_credential,
+                auth_discovery=auth_discovery,
+                tool_config=tool_config,
+            )
+        else:
+            # Base MCPToolset doesn't support tool_config parameter
+            super().__init__(
+                connection_params=connection_params,
+                tool_filter=tool_filter,
+                auth_scheme=auth_scheme,
+                auth_credential=auth_credential,
+            )
+
+        self._tool_cache = []
+        self._credential_manager = credential_manager
+
+    async def get_tools(self, readonly_context=None) -> list[MCPTool]:
         """
         Override get_tools to return EmbedResolvingMCPTool instances.
         """
+
+        if self._tool_cache:
+            return self._tool_cache
+
         # Get the original tools from the parent class
         original_tools = await super().get_tools(readonly_context)
 
@@ -294,7 +395,10 @@ class EmbedResolvingMCPToolset(MCPToolset):
             embed_resolving_tool = EmbedResolvingMCPTool(
                 original_mcp_tool=tool,
                 tool_config=tool_specific_config,
+                credential_manager=self._credential_manager,
             )
             embed_resolving_tools.append(embed_resolving_tool)
 
+        self._tool_cache = embed_resolving_tools
         return embed_resolving_tools
+

solace_agent_mesh/agent/adk/models/lite_llm.py

@@ -699,6 +699,11 @@ class LiteLlm(BaseLlm):
         super().__init__(model=model, **kwargs)
         self._additional_args = kwargs.copy()
 
+        # Remove handlers added by LiteLLM as they produce duplicate and misformatted logs.
+        # Logging is an application concern and libraries should not set handlers/formatters.
+        for logger_name in ["LiteLLM", "LiteLLM Proxy", "LiteLLM Router", "litellm"]:
+            logging.getLogger(logger_name).handlers.clear()
+
         # Validate and store cache strategy
         valid_strategies = ["none", "5m", "1h"]
         if cache_strategy not in valid_strategies:

solace_agent_mesh/agent/adk/runner.py

@@ -38,6 +38,7 @@ async def run_adk_async_task_thread_wrapper(
     adk_content: adk_types.Content,
     run_config: RunConfig,
     a2a_context: dict[str, Any],
+    append_context_event: bool = True,
 ):
     """
     Wrapper to run the async ADK task.
@@ -73,9 +74,8 @@ async def run_adk_async_task_thread_wrapper(
             logical_task_id,
         )
 
-        if adk_session and component.session_service:
+        if adk_session and component.session_service and append_context_event:
             context_setting_invocation_id = logical_task_id
-            original_message = a2a_context.pop("original_solace_message", None)
             try:
                 context_setting_event = ADKEvent(
                     invocation_id=context_setting_invocation_id,
@@ -107,15 +107,13 @@ async def run_adk_async_task_thread_wrapper(
                     e_append,
                     exc_info=True,
                 )
-            finally:
-                if original_message:
-                    a2a_context["original_solace_message"] = original_message
         else:
-            log.warning(
-                "%s Could not inject a2a_context into ADK session state via event for task %s (session or session_service invalid). Tool scope filtering might not work.",
-                component.log_identifier,
-                logical_task_id,
-            )
+            if append_context_event:
+                log.warning(
+                    "%s Could not inject a2a_context into ADK session state via event for task %s (session or session_service invalid). Tool scope filtering might not work.",
+                    component.log_identifier,
+                    logical_task_id,
+                )
 
         is_paused = await run_adk_async_task(
             component,
@@ -198,7 +196,7 @@ async def run_adk_async_task_thread_wrapper(
             "%s Bad Request for task %s: %s.",
             component.log_identifier,
             logical_task_id,
-            e.message
+            e.message,
         )
         raise
     except Exception as e:
@@ -345,7 +343,7 @@ async def run_adk_async_task(
             "%s Bad Request for task %s: %s.",
             component.log_identifier,
             logical_task_id,
-            e.message
+            e.message,
         )
         raise
     except Exception as e:

solace_agent_mesh/agent/adk/services.py

@@ -5,28 +5,32 @@ Initializes ADK Services based on configuration.
 import logging
 import os
 import re
-from typing import Dict, Optional, List, Any
-from typing_extensions import override
-
-from google.genai import types as adk_types
+from typing import Any, Dict, List, Optional
 
-
-from google.adk.sessions import (
-    BaseSessionService,
-    InMemorySessionService,
-    DatabaseSessionService,
-    VertexAiSessionService,
-)
 from google.adk.artifacts import (
     BaseArtifactService,
-    InMemoryArtifactService,
     GcsArtifactService,
+    InMemoryArtifactService,
+)
+from google.adk.auth.credential_service.base_credential_service import (
+    BaseCredentialService,
+)
+from google.adk.auth.credential_service.in_memory_credential_service import (
+    InMemoryCredentialService,
 )
 from google.adk.memory import (
     BaseMemoryService,
     InMemoryMemoryService,
     VertexAiRagMemoryService,
 )
+from google.adk.sessions import (
+    BaseSessionService,
+    DatabaseSessionService,
+    InMemorySessionService,
+    VertexAiSessionService,
+)
+from google.genai import types as adk_types
+from typing_extensions import override
 
 from .artifacts.filesystem_artifact_service import FilesystemArtifactService
 
@@ -170,9 +174,9 @@ def initialize_session_service(component) -> BaseSessionService:
     config = component.get_config("session_service", {})
 
     # Handle both dict and SessionServiceConfig object
-    if hasattr(config, 'type'):
+    if hasattr(config, "type"):
         service_type = config.type.lower()
-        db_url = getattr(config, 'database_url', None)
+        db_url = getattr(config, "database_url", None)
     else:
         service_type = config.get("type", "memory").lower()
         db_url = config.get("database_url")
@@ -390,3 +394,35 @@ def initialize_memory_service(component) -> BaseMemoryService:
         raise ValueError(
             f"{component.log_identifier} Unsupported memory service type: {service_type}"
         )
+
+
+def initialize_credential_service(component) -> BaseCredentialService | None:
+    """Initializes the ADK Credential Service based on configuration."""
+    config = component.get_config("credential_service", None)
+
+    # If no credential service is configured, return None
+    if config is None:
+        log.info(
+            "%s No credential service configured, skipping initialization",
+            component.log_identifier,
+        )
+        return None
+
+    # Handle both dict and CredentialServiceConfig object
+    if hasattr(config, "type"):
+        service_type = config.type.lower()
+    else:
+        service_type = config.get("type", "memory").lower()
+
+    log.info(
+        "%s Initializing Credential Service of type: %s",
+        component.log_identifier,
+        service_type,
+    )
+
+    if service_type == "memory":
+        return InMemoryCredentialService()
+    else:
+        raise ValueError(
+            f"{component.log_identifier} Unsupported credential service type: {service_type}"
+        )

solace_agent_mesh/agent/adk/setup.py

@@ -2,39 +2,31 @@
 Handles ADK Agent and Runner initialization, including tool loading and callback assignment.
 """
 
-import logging
-from typing import Dict, List, Optional, Union, Callable, Tuple, Set, Any, TYPE_CHECKING, Type
 import functools
 import inspect
-from solace_ai_connector.common.utils import import_module
-from ...common.utils.type_utils import is_subclass_by_name
-
+import logging
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple, Type, Union
 
-from .app_llm_agent import AppLlmAgent
-from .tool_wrapper import ADKToolWrapper
-from .embed_resolving_mcp_toolset import EmbedResolvingMCPToolset
-from google.adk.runners import Runner
-from google.adk.models import BaseLlm
-from google.adk.tools import BaseTool, ToolContext
 from google.adk import tools as adk_tools_module
 from google.adk.agents.callback_context import CallbackContext
+from google.adk.models import BaseLlm
 from google.adk.models.llm_request import LlmRequest
 from google.adk.models.llm_response import LlmResponse
+from google.adk.runners import Runner
+from google.adk.tools import BaseTool, ToolContext
 from google.adk.tools.mcp_tool.mcp_session_manager import (
     SseServerParams,
     StdioConnectionParams,
     StreamableHTTPServerParams,
-
 )
-
 from mcp import StdioServerParameters
+from solace_ai_connector.common.utils import import_module
 
-if TYPE_CHECKING:
-    from ..sac.component import SamAgentComponent
-
-from ..tools.registry import tool_registry
-from ..tools.tool_definition import BuiltinTool
+from ...agent.adk import callbacks as adk_callbacks
+from ...agent.adk.models.lite_llm import LiteLlm
+from ...common.utils.type_utils import is_subclass_by_name
 from ..tools.dynamic_tool import DynamicTool, DynamicToolProvider
+from ..tools.registry import tool_registry
 from ..tools.tool_config_types import (
     AnyToolConfig,
     BuiltinToolConfig,
@@ -42,10 +34,13 @@ from ..tools.tool_config_types import (
     McpToolConfig,
     PythonToolConfig,
 )
+from ..tools.tool_definition import BuiltinTool
+from .app_llm_agent import AppLlmAgent
+from .embed_resolving_mcp_toolset import EmbedResolvingMCPToolset
+from .tool_wrapper import ADKToolWrapper
 
-
-from ...agent.adk import callbacks as adk_callbacks
-from ...agent.adk.models.lite_llm import LiteLlm
+if TYPE_CHECKING:
+    from ..sac.component import SamAgentComponent
 
 log = logging.getLogger(__name__)
 
@@ -585,11 +580,43 @@ async def _load_mcp_tool(component: "SamAgentComponent", tool_config: Dict) -> T
             tool_config_model.tool_name,
         )
 
-    mcp_toolset_instance = EmbedResolvingMCPToolset(
-        connection_params=connection_params,
-        tool_filter=tool_filter_list,
-        tool_config=tool_config,
-    )
+    additional_params = {}
+    try:
+        from solace_agent_mesh_enterprise.auth.tool_configurator import (
+            configure_mcp_tool,
+        )
+
+        try:
+            # Call the tool configurator with MCP-specific context
+            additional_params = configure_mcp_tool(
+                tool_type="mcp",
+                tool_config=tool_config,
+                connection_params=connection_params,
+                tool_filter=tool_filter_list,
+            )
+        except Exception as e:
+            log.error(
+                "%s Tool configurator failed for %s: %s",
+                component.log_identifier,
+                tool_config.get("name", "unknown"),
+                e,
+            )
+            # Continue with normal tool creation if configurator fails
+            additional_params = {}
+    except ImportError:
+        pass
+
+    # Create the EmbedResolvingMCPToolset with base parameters
+    toolset_params = {
+        "connection_params": connection_params,
+        "tool_filter": tool_filter_list,
+        "tool_config": tool_config,
+    }
+
+    # Merge additional parameters from configurator
+    toolset_params.update(additional_params)
+
+    mcp_toolset_instance = EmbedResolvingMCPToolset(**toolset_params)
     mcp_toolset_instance.origin = "mcp"
 
     log.info(
@@ -891,7 +918,7 @@ def initialize_adk_agent(
         callbacks_in_order_for_before_model.append(
             adk_callbacks.repair_history_callback
         )
-        log.info(
+        log.debug(
             "%s Added repair_history_callback to before_model chain.",
             component.log_identifier,
         )
@@ -900,7 +927,7 @@ def initialize_adk_agent(
             callbacks_in_order_for_before_model.append(
                 component._inject_peer_tools_callback
             )
-            log.info(
+            log.debug(
                 "%s Added _inject_peer_tools_callback to before_model chain.",
                 component.log_identifier,
             )
@@ -909,7 +936,7 @@ def initialize_adk_agent(
             callbacks_in_order_for_before_model.append(
                 component._filter_tools_by_capability_callback
             )
-            log.info(
+            log.debug(
                 "%s Added _filter_tools_by_capability_callback to before_model chain.",
                 component.log_identifier,
             )
@@ -917,7 +944,7 @@ def initialize_adk_agent(
             callbacks_in_order_for_before_model.append(
                 component._inject_gateway_instructions_callback
             )
-            log.info(
+            log.debug(
                 "%s Added _inject_gateway_instructions_callback to before_model chain.",
                 component.log_identifier,
             )
@@ -930,7 +957,7 @@ def initialize_adk_agent(
         callbacks_in_order_for_before_model.append(
             dynamic_instruction_callback_with_component
         )
-        log.info(
+        log.debug(
             "%s Added inject_dynamic_instructions_callback to before_model chain.",
             component.log_identifier,
         )
@@ -957,7 +984,7 @@ def initialize_adk_agent(
             return None
 
         agent.before_model_callback = final_before_model_wrapper
-        log.info(
+        log.debug(
             "%s Final before_model_callback chain (Solace logging now occurs last) assigned to agent.",
             component.log_identifier,
         )
@@ -967,7 +994,7 @@ def initialize_adk_agent(
             host_component=component,
         )
         agent.before_tool_callback = tool_invocation_start_cb_with_component
-        log.info(
+        log.debug(
             "%s Assigned notify_tool_invocation_start_callback as before_tool_callback.",
             component.log_identifier,
         )
@@ -1053,7 +1080,7 @@ def initialize_adk_agent(
                 return tool_response
 
         agent.after_tool_callback = chained_after_tool_callback
-        log.info(
+        log.debug(
             "%s Chained 'manage_large_mcp_tool_responses_callback' and 'after_tool_callback_inject_metadata' as after_tool_callback.",
             component.log_identifier,
         )
@@ -1067,7 +1094,7 @@ def initialize_adk_agent(
             adk_callbacks.process_artifact_blocks_callback, host_component=component
         )
         callbacks_in_order_for_after_model.append(artifact_block_cb)
-        log.info(
+        log.debug(
             "%s Added process_artifact_blocks_callback to after_model chain.",
             component.log_identifier,
         )
@@ -1077,7 +1104,7 @@ def initialize_adk_agent(
             adk_callbacks.auto_continue_on_max_tokens_callback, host_component=component
         )
         callbacks_in_order_for_after_model.append(auto_continue_cb)
-        log.info(
+        log.debug(
             "%s Added auto_continue_on_max_tokens_callback to after_model chain.",
             component.log_identifier,
         )
@@ -1110,7 +1137,7 @@ def initialize_adk_agent(
             return None
 
         agent.after_model_callback = final_after_model_wrapper
-        log.info(
+        log.debug(
             "%s Chained all after_model_callbacks and assigned to agent.",
             component.log_identifier,
         )
@@ -1157,9 +1184,10 @@ def initialize_adk_runner(component) -> Runner:
             session_service=component.session_service,
             artifact_service=component.artifact_service,
             memory_service=component.memory_service,
+            credential_service=component.credential_service,
         )
         log.info("%s ADK Runner created successfully.", component.log_identifier)
         return runner
     except Exception as e:
         log.error("%s Failed to create ADK Runner: %s", component.log_identifier, e)
-        raise
+        raise