From e96cee313baa9b70f937885735c990bf7bd655c3 Mon Sep 17 00:00:00 2001 From: googleSandy Date: Wed, 10 Sep 2025 17:45:28 -0700 Subject: [PATCH 1/3] Update requirements.txt --- run-with-google-adk/requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/run-with-google-adk/requirements.txt b/run-with-google-adk/requirements.txt index 1503a5bb..e8f51ee4 100644 --- a/run-with-google-adk/requirements.txt +++ b/run-with-google-adk/requirements.txt @@ -1,6 +1,6 @@ google-cloud-aiplatform==1.97.0 markdown uv -google-adk[eval]==1.3.0 +google-adk[eval]==1.14.0 google-genai==1.20.0 pandas From e297a29a3953953b819bee087885c3d04baeabc9 Mon Sep 17 00:00:00 2001 From: googleSandy Date: Wed, 10 Sep 2025 17:47:54 -0700 Subject: [PATCH 2/3] Update agent.py Update and cleaned up agent.py to work with latest ADK 1.14.0 --- .../google_mcp_security_agent/agent.py | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) diff --git a/run-with-google-adk/google_mcp_security_agent/agent.py b/run-with-google-adk/google_mcp_security_agent/agent.py index 06782a6c..8f98ae06 100644 --- a/run-with-google-adk/google_mcp_security_agent/agent.py +++ b/run-with-google-adk/google_mcp_security_agent/agent.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from google.adk.agents.llm_agent import LlmAgent -from google.adk.tools.mcp_tool.mcp_toolset import StdioServerParameters, StdioConnectionParams +from google.adk.tools.mcp_tool.mcp_toolset import StdioServerParameters import os import logging @@ -62,8 +62,7 @@ def get_all_tools(): if os.environ.get("LOAD_SCC_MCP") == "Y": scc_tools = MCPToolSetWithSchemaAccess( - connection_params=StdioConnectionParams( - server_params=StdioServerParameters( + connection_params=StdioServerParameters( command='uv', args=[ "--directory", uv_dir_prefix + "/scc", @@ -71,15 +70,13 @@ def get_all_tools(): "scc_mcp.py" ] ), - timeout=timeout), tool_set_name="scc", errlog=errlog_ae ) if os.environ.get("LOAD_SECOPS_MCP") == "Y": secops_tools = MCPToolSetWithSchemaAccess( - connection_params=StdioConnectionParams( - server_params=StdioServerParameters( + connection_params=StdioServerParameters( command='uv', args=[ "--directory", uv_dir_prefix + "/secops/secops_mcp", @@ -89,15 +86,13 @@ def get_all_tools(): "server.py" ] ), - timeout=timeout), tool_set_name="secops_mcp", errlog=errlog_ae ) if os.environ.get("LOAD_GTI_MCP") == "Y": gti_tools = MCPToolSetWithSchemaAccess( - connection_params=StdioConnectionParams( - server_params=StdioServerParameters( + connection_params=StdioServerParameters( command='uv', args=[ "--directory", uv_dir_prefix + "/gti/gti_mcp", @@ -107,7 +102,6 @@ def get_all_tools(): "server.py" ] ), - timeout=timeout), tool_set_name="gti_mcp", errlog=errlog_ae ) @@ -115,8 +109,7 @@ def get_all_tools(): if os.environ.get("LOAD_SECOPS_SOAR_MCP") == "Y": secops_soar_tools = MCPToolSetWithSchemaAccess( - connection_params=StdioConnectionParams( - server_params=StdioServerParameters( + connection_params=StdioServerParameters( command='uv', args=[ "--directory", uv_dir_prefix + "/secops-soar/secops_soar_mcp", @@ -128,7 +121,6 @@ def get_all_tools(): os.environ.get("SECOPS_INTEGRATIONS","CSV,OKTA") ] ), - timeout=timeout), tool_set_name="secops_soar_mcp", errlog=errlog_ae ) From 5de0ab9ed6f8612381005d9273636cb25aea05dc Mon Sep 17 00:00:00 2001 From: googleSandy Date: Wed, 10 Sep 2025 17:50:17 -0700 Subject: [PATCH 3/3] Update extensions.py --- .../utils_extensions_cbs_tools/extensions.py | 51 +++++-------------- 1 file changed, 14 insertions(+), 37 deletions(-) diff --git a/run-with-google-adk/utils_extensions_cbs_tools/extensions.py b/run-with-google-adk/utils_extensions_cbs_tools/extensions.py index 79fb3324..bf8bf512 100644 --- a/run-with-google-adk/utils_extensions_cbs_tools/extensions.py +++ b/run-with-google-adk/utils_extensions_cbs_tools/extensions.py @@ -13,17 +13,13 @@ # limitations under the License. -# imports for overriding `get_tools` -from typing_extensions import override -from google.adk.tools.mcp_tool.mcp_session_manager import retry_on_closed_resource from typing import List from typing import Optional, Union, TextIO from google.adk.agents.readonly_context import ReadonlyContext -from google.adk.tools.mcp_tool.mcp_tool import MCPTool, BaseTool -from google.adk.tools.mcp_tool.mcp_session_manager import StdioServerParameters, StdioConnectionParams, SseConnectionParams,StreamableHTTPConnectionParams -from mcp.types import ListToolsResult +from google.adk.tools.base_tool import BaseTool +from google.adk.tools.mcp_tool.mcp_toolset import StdioServerParameters from .cache import tools_cache -from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset, ToolPredicate +from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset import sys import logging @@ -42,31 +38,23 @@ def __init__( self, *, tool_set_name: str, # <-- new parameter - connection_params: Union[ - StdioServerParameters, - StdioConnectionParams, - SseConnectionParams, - StreamableHTTPConnectionParams, - ], - tool_filter: Optional[Union[ToolPredicate, List[str]]] = None, + connection_params: StdioServerParameters, + tool_filter: Optional[List[str]] = None, errlog: TextIO = sys.stderr, ): super().__init__( connection_params=connection_params, - tool_filter=tool_filter, errlog=errlog ) self.tool_set_name = tool_set_name logging.info(f"MCPToolSetWithSchemaAccess initialized with tool_set_name: '{self.tool_set_name}'") self._session = None - @retry_on_closed_resource("_reinitialize_session") - @override async def get_tools( self, readonly_context: Optional[ReadonlyContext] = None, ) -> List[BaseTool]: - """Return all tools in the toolset based on the provided context. + """Return all tools in the toolset based on the provided context with caching. Args: readonly_context: Context used to filter tools available to the agent. @@ -75,28 +63,17 @@ async def get_tools( Returns: List[BaseTool]: A list of tools available under the specified context. """ - # Get session from session manager - if not self._session: - self._session = await self._mcp_session_manager.create_session() - + # Check cache first if self.tool_set_name in tools_cache.keys(): logging.info(f"Tools found in cache for toolset {self.tool_set_name}, returning them") return tools_cache[self.tool_set_name] else: - logging.info(f"No tools found in cache for toolset {self.tool_set_name}, loading") - - tools_response: ListToolsResult = await self._session.list_tools() - - # Apply filtering based on context and tool_filter - tools = [] - for tool in tools_response.tools: - mcp_tool = MCPTool( - mcp_tool=tool, - mcp_session_manager=self._mcp_session_manager, - ) - - if self._is_tool_selected(mcp_tool, readonly_context): - tools.append(mcp_tool) + logging.info(f"No tools found in cache for toolset {self.tool_set_name}, loading from parent") + # Get tools from parent class + tools = await super().get_tools(readonly_context) + + # Cache the tools tools_cache[self.tool_set_name] = tools - return tools \ No newline at end of file + logging.info(f"Cached {len(tools)} tools for toolset {self.tool_set_name}") + return tools