fix(code_executors): stop ContainerCodeExecutor leaking via atexit

src/google/adk/code_executors/container_code_executor.py

@@ -18,6 +18,7 @@
 import logging
 import os
 from typing import Optional
+import weakref
 
 import docker
 from docker.client import DockerClient
@@ -139,8 +140,15 @@ def __init__(
     # Initialize the container.
     self.__init_container()
 
-    # Close the container when the on exit.
-    atexit.register(self.__cleanup_container)
+    # Close the container on interpreter exit. Register a weakref.proxy so the
+    # process-global atexit registry does not keep a strong reference to this
+    # executor: otherwise every ContainerCodeExecutor ever created (and its
+    # long-lived Docker container) would be retained until interpreter exit,
+    # even after the executor is no longer used. A server that builds executors
+    # per app/agent/session would leak them unboundedly.
+    atexit.register(
+        ContainerCodeExecutor.__cleanup_container, weakref.proxy(self)
+    )
 
   @override
   def execute_code(
@@ -219,12 +227,23 @@ def __init_container(self):
     # Verify the container is able to run python3.
     self._verify_python_installation()
 
-  def __cleanup_container(self):
-    """Closes the container on exit."""
-    if not self._container:
+  @staticmethod
+  def __cleanup_container(self: ContainerCodeExecutor):
+    """Closes the container on exit.
+
+    Registered with ``atexit`` via a ``weakref.proxy`` so it does not keep the
+    executor alive. If the executor has already been garbage collected the
+    proxy is dead and there is nothing left to clean up.
+    """
+    try:
+      container = self._container
+    except ReferenceError:
+      return
+
+    if not container:
       return
 
     logger.info('[Cleanup] Stopping the container...')
-    self._container.stop()
-    self._container.remove()
-    logger.info('Container %s stopped and removed.', self._container.id)
+    container.stop()
+    container.remove()
+    logger.info('Container %s stopped and removed.', container.id)

tests/unittests/code_executors/test_container_code_executor.py

@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for the ContainerCodeExecutor container hardening defaults."""
+"""Tests for the ContainerCodeExecutor container hardening and lifecycle."""
 
+import gc
 from unittest import mock
+import weakref
 
 from google.adk.code_executors.container_code_executor import ContainerCodeExecutor
 
@@ -56,3 +58,41 @@ def test_container_network_can_be_explicitly_enabled(mock_docker):
 
   _, kwargs = client.containers.run.call_args
   assert not kwargs['network_disabled']
+
+
+@mock.patch('google.adk.code_executors.container_code_executor.docker')
+def test_executor_is_not_retained_by_atexit(mock_docker):
+  """The atexit handler must not keep the executor (or its container) alive.
+
+  The exit handler is registered via a ``weakref.proxy`` so that a discarded
+  executor can be garbage collected instead of surviving (together with its
+  running Docker container) until interpreter exit. Registering a bound method
+  would keep a strong reference in the process-global atexit registry and leak
+  every executor ever created.
+  """
+  client = _mock_docker_client()
+  mock_docker.from_env.return_value = client
+
+  executor = ContainerCodeExecutor(image='test-image')
+  ref = weakref.ref(executor)
+
+  del executor
+  gc.collect()
+
+  assert ref() is None
+
+
+@mock.patch('google.adk.code_executors.container_code_executor.docker')
+def test_cleanup_stops_and_removes_container(mock_docker):
+  """The exit handler still stops and removes a live executor's container."""
+  client = _mock_docker_client()
+  container = client.containers.run.return_value
+  mock_docker.from_env.return_value = client
+
+  executor = ContainerCodeExecutor(image='test-image')
+  # Name-mangled staticmethod registered with atexit; invoke it directly to
+  # emulate interpreter exit while the executor is still referenced.
+  ContainerCodeExecutor._ContainerCodeExecutor__cleanup_container(executor)
+
+  container.stop.assert_called_once()
+  container.remove.assert_called_once()