feat: add configurable resource limits for subprocesses in BashTool

PiperOrigin-RevId: 894338571

Author: shukladivyansh

src/google/adk/tools/bash_tool.py

@@ -21,6 +21,7 @@
 import logging
 import os
 import pathlib
+import resource
 import shlex
 import signal
 from typing import Any
@@ -32,16 +33,25 @@
 from .base_tool import BaseTool
 from .tool_context import ToolContext
 
+logger = logging.getLogger("google_adk." + __name__)
+
 
 @dataclasses.dataclass(frozen=True)
 class BashToolPolicy:
-  """Configuration for allowed bash commands based on prefix matching.
+  """Configuration for allowed bash commands and resource limits.
 
   Set allowed_command_prefixes to ("*",) to allow all commands (default),
   or explicitly list allowed prefixes.
+
+  Values for max_memory_bytes, max_file_size_bytes, and max_child_processes
+  will be enforced upon the spawned subprocess.
   """
 
   allowed_command_prefixes: tuple[str, ...] = ("*",)
+  timeout_seconds: Optional[int] = 30
+  max_memory_bytes: Optional[int] = None
+  max_file_size_bytes: Optional[int] = None
+  max_child_processes: Optional[int] = None
 
 
 def _validate_command(command: str, policy: BashToolPolicy) -> Optional[str]:
@@ -61,6 +71,29 @@ def _validate_command(command: str, policy: BashToolPolicy) -> Optional[str]:
   return f"Command blocked. Permitted prefixes are: {allowed}"
 
 
+def _set_resource_limits(policy: BashToolPolicy) -> None:
+  """Sets resource limits for the subprocess based on the provided policy."""
+  try:
+    resource.setrlimit(resource.RLIMIT_CORE, (0, 0))
+    if policy.max_memory_bytes:
+      resource.setrlimit(
+          resource.RLIMIT_AS,
+          (policy.max_memory_bytes, policy.max_memory_bytes),
+      )
+    if policy.max_file_size_bytes:
+      resource.setrlimit(
+          resource.RLIMIT_FSIZE,
+          (policy.max_file_size_bytes, policy.max_file_size_bytes),
+      )
+    if policy.max_child_processes:
+      resource.setrlimit(
+          resource.RLIMIT_NPROC,
+          (policy.max_child_processes, policy.max_child_processes),
+      )
+  except (ValueError, OSError) as e:
+    logger.warning("Failed to set resource limits: %s", e)
+
+
 @features.experimental(features.FeatureName.SKILL_TOOLSET)
 class ExecuteBashTool(BaseTool):
   """Tool to execute a validated bash command within a workspace directory."""
@@ -144,20 +177,25 @@ async def run_async(
           stdout=asyncio.subprocess.PIPE,
           stderr=asyncio.subprocess.PIPE,
           start_new_session=True,
+          preexec_fn=lambda: _set_resource_limits(self._policy),
       )
 
       try:
         stdout, stderr = await asyncio.wait_for(
-            process.communicate(), timeout=30
+            process.communicate(), timeout=self._policy.timeout_seconds
         )
       except asyncio.TimeoutError:
         try:
-          os.killpg(process.pid, signal.SIGKILL)
+          if process.pid:
+            os.killpg(process.pid, signal.SIGKILL)
         except ProcessLookupError:
           pass
         stdout, stderr = await process.communicate()
         return {
-            "error": "Command timed out after 30 seconds.",
+            "error": (
+                f"Command timed out after {self._policy.timeout_seconds}"
+                " seconds."
+            ),
             "stdout": (
                 stdout.decode(errors="replace")
                 if stdout
@@ -176,7 +214,6 @@ async def run_async(
             os.killpg(process.pid, signal.SIGKILL)
         except ProcessLookupError:
           pass
-
       return {
           "stdout": (
               stdout.decode(errors="replace")
@@ -191,7 +228,6 @@ async def run_async(
           "returncode": process.returncode,
       }
     except Exception as e:  # pylint: disable=broad-except
-      logger = logging.getLogger("google_adk." + __name__)
       logger.exception("ExecuteBashTool execution failed")
 
       stdout_res = (

tests/unittests/tools/test_bash_tool.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import asyncio
+import resource
 import signal
 from unittest import mock
 
@@ -241,3 +242,34 @@ async def test_no_command(self, workspace, tool_context_confirmed):
     result = await tool.run_async(args={}, tool_context=tool_context_confirmed)
     assert "error" in result
     assert "required" in result["error"].lower()
+
+  @pytest.mark.asyncio
+  async def test_resource_limits_set(self, workspace, tool_context_confirmed):
+    policy = bash_tool.BashToolPolicy(
+        max_memory_bytes=100 * 1024 * 1024,
+        max_file_size_bytes=50 * 1024 * 1024,
+        max_child_processes=10,
+    )
+    tool = bash_tool.ExecuteBashTool(workspace=workspace, policy=policy)
+    mock_process = mock.AsyncMock()
+    mock_process.communicate.return_value = (b"", b"")
+    mock_exec = mock.AsyncMock(return_value=mock_process)
+
+    with mock.patch("asyncio.create_subprocess_exec", mock_exec):
+      await tool.run_async(
+          args={"command": "ls"},
+          tool_context=tool_context_confirmed,
+      )
+      assert "preexec_fn" in mock_exec.call_args.kwargs
+      preexec_fn = mock_exec.call_args.kwargs["preexec_fn"]
+
+      mock_setrlimit = mock.create_autospec(resource.setrlimit, instance=True)
+      with mock.patch("resource.setrlimit", mock_setrlimit):
+        preexec_fn()
+        mock_setrlimit.assert_any_call(resource.RLIMIT_CORE, (0, 0))
+        mock_setrlimit.assert_any_call(
+            resource.RLIMIT_AS, (100 * 1024 * 1024, 100 * 1024 * 1024)
+        )
+        mock_setrlimit.assert_any_call(
+            resource.RLIMIT_FSIZE, (50 * 1024 * 1024, 50 * 1024 * 1024)
+        )