Skip to content

Commit 161742f

Browse files
Merge branch 'main' into fix/parallel-agent-pre-311-aclose-error
2 parents 33088ee + 7ae83b2 commit 161742f

12 files changed

Lines changed: 814 additions & 227 deletions

File tree

src/google/adk/agents/config_agent_utils.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,31 @@ def _resolve_agent_class(agent_class: str) -> type[BaseAgent]:
8080
)
8181

8282

83+
_BLOCKED_YAML_KEYS = frozenset({"args"})
84+
_ENFORCE_DENYLIST = False
85+
86+
87+
def _set_enforce_denylist(value: bool) -> None:
88+
global _ENFORCE_DENYLIST
89+
_ENFORCE_DENYLIST = value
90+
91+
92+
def _check_config_for_blocked_keys(node: Any, filename: str) -> None:
93+
"""Recursively check if the configuration contains any blocked keys."""
94+
if isinstance(node, dict):
95+
for key, value in node.items():
96+
if key in _BLOCKED_YAML_KEYS:
97+
raise ValueError(
98+
f"Blocked key {key!r} found in {filename!r}. "
99+
f"The '{key}' field is not allowed in agent configurations "
100+
"because it can execute arbitrary code."
101+
)
102+
_check_config_for_blocked_keys(value, filename)
103+
elif isinstance(node, list):
104+
for item in node:
105+
_check_config_for_blocked_keys(item, filename)
106+
107+
83108
def _load_config_from_path(config_path: str) -> AgentConfig:
84109
"""Load an agent's configuration from a YAML file.
85110
@@ -100,6 +125,9 @@ def _load_config_from_path(config_path: str) -> AgentConfig:
100125
with open(config_path, "r", encoding="utf-8") as f:
101126
config_data = yaml.safe_load(f)
102127

128+
if _ENFORCE_DENYLIST:
129+
_check_config_for_blocked_keys(config_data, config_path)
130+
103131
return AgentConfig.model_validate(config_data)
104132

105133

src/google/adk/cli/fast_api.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,12 @@ def get_fast_api_app(
148148
The configured FastAPI application instance.
149149
"""
150150

151+
# Enable denylist enforcement for config loads if web UI is enabled.
152+
if web:
153+
from ..agents import config_agent_utils
154+
155+
config_agent_utils._set_enforce_denylist(True)
156+
151157
# Set up eval managers.
152158
if eval_storage_uri:
153159
gcs_eval_managers = evals.create_gcs_eval_managers_from_uri(

src/google/adk/tools/load_mcp_resource_tool.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,7 @@
2929
from .base_tool import BaseTool
3030

3131
if TYPE_CHECKING:
32-
from mcp_toolset import McpToolset
33-
32+
from .mcp_tool.mcp_toolset import McpToolset
3433
from .tool_context import ToolContext
3534

3635
logger = logging.getLogger("google_adk." + __name__)
@@ -39,7 +38,7 @@
3938
class LoadMcpResourceTool(BaseTool):
4039
"""A tool that loads the MCP resources and adds them to the session."""
4140

42-
def __init__(self, mcp_toolset: McpToolset):
41+
def __init__(self, mcp_toolset: McpToolset) -> None:
4342
super().__init__(
4443
name="load_mcp_resource",
4544
description="""Loads resources from the MCP server.

src/google/adk/tools/mcp_tool/__init__.py

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -17,17 +17,19 @@
1717
try:
1818
from .conversion_utils import adk_to_mcp_tool_type
1919
from .conversion_utils import gemini_to_json_schema
20-
from .mcp_session_manager import SseConnectionParams
21-
from .mcp_session_manager import StdioConnectionParams
22-
from .mcp_session_manager import StreamableHTTPConnectionParams
23-
from .mcp_tool import MCPTool
24-
from .mcp_tool import McpTool
25-
from .mcp_toolset import MCPToolset
26-
from .mcp_toolset import McpToolset
20+
from .mcp_session_manager import MCPSessionManager as MCPSessionManager
21+
from .mcp_session_manager import SseConnectionParams as SseConnectionParams
22+
from .mcp_session_manager import StdioConnectionParams as StdioConnectionParams
23+
from .mcp_session_manager import StreamableHTTPConnectionParams as StreamableHTTPConnectionParams
24+
from .mcp_tool import MCPTool as MCPTool
25+
from .mcp_tool import McpTool as McpTool
26+
from .mcp_toolset import MCPToolset as MCPToolset
27+
from .mcp_toolset import McpToolset as McpToolset
2728

2829
__all__.extend([
2930
'adk_to_mcp_tool_type',
3031
'gemini_to_json_schema',
32+
'MCPSessionManager',
3133
'McpTool',
3234
'MCPTool',
3335
'McpToolset',

0 commit comments

Comments
 (0)