From c37950b4f228ec3960b60aca4e343030d2defc0d Mon Sep 17 00:00:00 2001 From: Pedro Melendez Date: Thu, 13 Nov 2025 12:34:16 -0800 Subject: [PATCH 1/5] Fix deepcopy recursion in VertexAiCodeExecutor #1762 and #1730 --- .../code_executors/vertex_ai_code_executor.py | 29 +-- .../test_vertex_ai_code_executor.py | 182 ++++++++++++++++++ 2 files changed, 200 insertions(+), 11 deletions(-) create mode 100644 tests/unittests/code_executors/test_vertex_ai_code_executor.py diff --git a/src/google/adk/code_executors/vertex_ai_code_executor.py b/src/google/adk/code_executors/vertex_ai_code_executor.py index a6a0ec8eb5..2992d0964d 100644 --- a/src/google/adk/code_executors/vertex_ai_code_executor.py +++ b/src/google/adk/code_executors/vertex_ai_code_executor.py @@ -20,6 +20,7 @@ from typing import Any from typing import Optional +from pydantic import PrivateAttr from typing_extensions import override from ..agents.invocation_context import InvocationContext @@ -113,14 +114,10 @@ class VertexAiCodeExecutor(BaseCodeExecutor): projects/123/locations/us-central1/extensions/456 """ - resource_name: str = None - """ - If set, load the existing resource name of the code interpreter extension - instead of creating a new one. - Format: projects/123/locations/us-central1/extensions/456 - """ + resource_name: Optional[str] = None - _code_interpreter_extension: Extension + # Mark as PrivateAttr so Pydantic and copy.deepcopy ignore it during validation/cloning + _code_interpreter_extension: Any = PrivateAttr(default=None) def __init__( self, @@ -137,9 +134,17 @@ def __init__( """ super().__init__(**data) self.resource_name = resource_name - self._code_interpreter_extension = _get_code_interpreter_extension( - self.resource_name - ) + # Note: We do not initialize _code_interpreter_extension here to avoid recursion errors + # during deployment deepcopy. It is initialized lazily in the property below. + + @property + def extension_client(self): + """Lazy loads the Vertex AI Extension client.""" + if self._code_interpreter_extension is None: + self._code_interpreter_extension = _get_code_interpreter_extension( + self.resource_name + ) + return self._code_interpreter_extension @override def execute_code( @@ -220,7 +225,9 @@ def _execute_code_interpreter( ] if session_id: operation_params['session_id'] = session_id - response = self._code_interpreter_extension.execute( + + # Use the lazy-loaded client property + response = self.extension_client.execute( operation_id='execute', operation_params=operation_params, ) diff --git a/tests/unittests/code_executors/test_vertex_ai_code_executor.py b/tests/unittests/code_executors/test_vertex_ai_code_executor.py new file mode 100644 index 0000000000..9a25f5fc59 --- /dev/null +++ b/tests/unittests/code_executors/test_vertex_ai_code_executor.py @@ -0,0 +1,182 @@ +import copy +from typing import List +import unittest +from unittest.mock import MagicMock +from unittest.mock import patch + +from google.adk.code_executors.vertex_ai_code_executor import CodeExecutionInput +from google.adk.code_executors.vertex_ai_code_executor import File +from google.adk.code_executors.vertex_ai_code_executor import VertexAiCodeExecutor + +InvocationContext = MagicMock + + +class TestVertexAiCodeExecutor(unittest.TestCase): + + def setUp(self): + """Set up common fixtures for the tests.""" + self.mock_resource_name = ( + 'projects/123/locations/us-central1/extensions/456' + ) + self.executor = VertexAiCodeExecutor(resource_name=self.mock_resource_name) + + def _create_mock_files(self, file_data: List[tuple]) -> List[File]: + """Helper to create File objects from (name, content, mime_type).""" + return [ + File(name=name, content=content, mime_type=mime_type) + for name, content, mime_type in file_data + ] + + # --- Test Initialization & Deepcopy Safety --- + + def test_init_is_lazy(self): + """Verifies __init__ does NOT create the external client.""" + self.assertIsNone(self.executor._code_interpreter_extension) + + def test_deepcopy_safety(self): + """Verifies that deepcopy works without RecursionError.""" + try: + executor_copy = copy.deepcopy(self.executor) + except RecursionError: + self.fail('deepcopy raised RecursionError! Lazy loading fix failed.') + + self.assertNotEqual(id(self.executor), id(executor_copy)) + + # --- Test Lazy Loading --- + + @patch('vertexai.preview.extensions.Extension') + def test_lazy_loading_and_caching(self, MockExtensionClass): + """Verifies client is created only on access and is cached.""" + + mock_client_instance = MockExtensionClass.return_value = MagicMock() + + # 1. Access the property to trigger instantiation (Lazy Loading) + with self.subTest(msg='Test Lazy Loading'): + client = self.executor.extension_client + MockExtensionClass.assert_called_once_with(self.mock_resource_name) + self.assertEqual(client, mock_client_instance) + + # 2. Access again to ensure no re-instantiation (Caching) + with self.subTest(msg='Test Caching'): + _ = self.executor.extension_client + MockExtensionClass.assert_called_once() + + # --- Test Execution Flow --- + + @patch('vertexai.preview.extensions.Extension') + def test_execute_code_flow(self, MockExtensionClass): + """Verifies execute_code correctly maps inputs, calls the client, and parses results.""" + + # 1. Setup Mocks and Response + mock_client = MagicMock() + MockExtensionClass.return_value = mock_client + MOCK_RESPONSE = { + 'execution_result': 'Final print output', + 'execution_error': '', + 'output_files': [ + {'name': 'plot.png', 'contents': 'base64_plot_string'}, + {'name': 'data.csv', 'contents': '1,2,3'}, + ], + } + mock_client.execute.return_value = MOCK_RESPONSE + + # 2. Input Data Preparation + input_files = self._create_mock_files( + [('input.txt', 'test content', 'text/plain')] + ) + input_data = CodeExecutionInput( + code='df.plot()', + execution_id='test-session-42', + input_files=input_files, + ) + context = MagicMock() + + # 3. Run execution + result = self.executor.execute_code(context, input_data) + + # 4. Verify client call arguments + _, kwargs = mock_client.execute.call_args + actual_code = kwargs['operation_params']['code'] + actual_files = kwargs['operation_params']['files'] + + # Assertions for dynamic parts + self.assertIn( + 'def explore_df(df: pd.DataFrame) -> None:', + actual_code, + 'Code payload must include the explore_df helper function.', + ) + self.assertTrue( + actual_code.strip().endswith('df.plot()'), + 'User code must be appended at the end of the payload.', + ) + self.assertNotIn( + 'mime_type', + actual_files[0], + "Files dict sent to client should NOT contain 'mime_type'.", + ) + + # Assertion for static parts + self.assertEqual(kwargs['operation_id'], 'execute') + self.assertEqual( + kwargs['operation_params']['session_id'], 'test-session-42' + ) + self.assertEqual( + kwargs['operation_params']['files'], + [ + # Ensure 'mime_type' is explicitly removed + {'name': 'input.txt', 'contents': 'test content'} + ], + ) + + # 5. Verify Output Parsing + self.assertEqual(result.stdout, MOCK_RESPONSE['execution_result']) + self.assertEqual(len(result.output_files), 2) + + with self.subTest(msg='Check Image File Parsing'): + image_file = result.output_files[0] + self.assertEqual(image_file.name, 'plot.png') + self.assertEqual(image_file.mime_type, 'image/png') + + with self.subTest(msg='Check CSV File Parsing'): + csv_file = result.output_files[1] + self.assertEqual(csv_file.name, 'data.csv') + self.assertEqual(csv_file.mime_type, 'text/csv') + + # --- Test Error Handling --- + + @patch('vertexai.preview.extensions.Extension') + def test_execute_code_api_exception(self, MockExtensionClass): + """Verifies that exceptions from the Vertex AI client bubble up correctly.""" + mock_client = MockExtensionClass.return_value = MagicMock() + + # Simulate a generic API failure (e.g. 500 error or Timeout) + mock_client.execute.side_effect = RuntimeError( + 'Vertex AI Service Unavailable' + ) + + input_data = CodeExecutionInput(code="print('fail')", input_files=[]) + context = MagicMock() + + # Verify the executor does not silently swallow critical errors + with self.assertRaises(RuntimeError) as cm: + self.executor.execute_code(context, input_data) + + self.assertEqual(str(cm.exception), 'Vertex AI Service Unavailable') + + @patch('vertexai.preview.extensions.Extension') + def test_execute_code_malformed_response(self, MockExtensionClass): + """Verifies behavior when API returns a response missing required keys.""" + mock_client = MockExtensionClass.return_value = MagicMock() + + # Simulate a response that lacks 'output_files' (contract violation) + mock_client.execute.return_value = { + 'execution_result': 'Success', + # 'output_files': [] <-- MISSING KEY + } + + input_data = CodeExecutionInput(code="print('ok')", input_files=[]) + context = MagicMock() + + # Expect a KeyError because the source code accesses ['output_files'] directly + with self.assertRaises(KeyError): + self.executor.execute_code(context, input_data) From 497b8b14714291c0cc8639da1cedbd5392ff9112 Mon Sep 17 00:00:00 2001 From: Pedro Melendez Date: Wed, 19 Nov 2025 16:22:26 -0800 Subject: [PATCH 2/5] Update src/google/adk/code_executors/vertex_ai_code_executor.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- src/google/adk/code_executors/vertex_ai_code_executor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/google/adk/code_executors/vertex_ai_code_executor.py b/src/google/adk/code_executors/vertex_ai_code_executor.py index 2992d0964d..925cc09def 100644 --- a/src/google/adk/code_executors/vertex_ai_code_executor.py +++ b/src/google/adk/code_executors/vertex_ai_code_executor.py @@ -117,7 +117,7 @@ class VertexAiCodeExecutor(BaseCodeExecutor): resource_name: Optional[str] = None # Mark as PrivateAttr so Pydantic and copy.deepcopy ignore it during validation/cloning - _code_interpreter_extension: Any = PrivateAttr(default=None) + _code_interpreter_extension: Optional["Extension"] = PrivateAttr(default=None) def __init__( self, From 49eeec10d7c27b23ab38f566c8f8f9f227bf8c1a Mon Sep 17 00:00:00 2001 From: Pedro Melendez Date: Wed, 19 Nov 2025 16:22:34 -0800 Subject: [PATCH 3/5] Update src/google/adk/code_executors/vertex_ai_code_executor.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- src/google/adk/code_executors/vertex_ai_code_executor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/google/adk/code_executors/vertex_ai_code_executor.py b/src/google/adk/code_executors/vertex_ai_code_executor.py index 925cc09def..c93ef85b47 100644 --- a/src/google/adk/code_executors/vertex_ai_code_executor.py +++ b/src/google/adk/code_executors/vertex_ai_code_executor.py @@ -138,7 +138,7 @@ def __init__( # during deployment deepcopy. It is initialized lazily in the property below. @property - def extension_client(self): + def extension_client(self) -> "Extension": """Lazy loads the Vertex AI Extension client.""" if self._code_interpreter_extension is None: self._code_interpreter_extension = _get_code_interpreter_extension( From fe3c6bf39d27c0850c22328f44503b57eb2a556a Mon Sep 17 00:00:00 2001 From: Pedro Melendez Date: Wed, 19 Nov 2025 16:23:07 -0800 Subject: [PATCH 4/5] Update tests/unittests/code_executors/test_vertex_ai_code_executor.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- tests/unittests/code_executors/test_vertex_ai_code_executor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unittests/code_executors/test_vertex_ai_code_executor.py b/tests/unittests/code_executors/test_vertex_ai_code_executor.py index 9a25f5fc59..3d58ace28c 100644 --- a/tests/unittests/code_executors/test_vertex_ai_code_executor.py +++ b/tests/unittests/code_executors/test_vertex_ai_code_executor.py @@ -20,7 +20,7 @@ def setUp(self): ) self.executor = VertexAiCodeExecutor(resource_name=self.mock_resource_name) - def _create_mock_files(self, file_data: List[tuple]) -> List[File]: + def _create_mock_files(self, file_data: List[tuple[str, str, str]]) -> List[File]: """Helper to create File objects from (name, content, mime_type).""" return [ File(name=name, content=content, mime_type=mime_type) From ac080acadd33a83233f62f39776c0e66aa2712c0 Mon Sep 17 00:00:00 2001 From: Pedro Melendez Date: Thu, 20 Nov 2025 16:52:51 -0800 Subject: [PATCH 5/5] Fixed formatting issues with pyink --- src/google/adk/code_executors/vertex_ai_code_executor.py | 4 ++-- .../unittests/code_executors/test_vertex_ai_code_executor.py | 4 +++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/google/adk/code_executors/vertex_ai_code_executor.py b/src/google/adk/code_executors/vertex_ai_code_executor.py index c93ef85b47..e14be59870 100644 --- a/src/google/adk/code_executors/vertex_ai_code_executor.py +++ b/src/google/adk/code_executors/vertex_ai_code_executor.py @@ -117,7 +117,7 @@ class VertexAiCodeExecutor(BaseCodeExecutor): resource_name: Optional[str] = None # Mark as PrivateAttr so Pydantic and copy.deepcopy ignore it during validation/cloning - _code_interpreter_extension: Optional["Extension"] = PrivateAttr(default=None) + _code_interpreter_extension: Optional['Extension'] = PrivateAttr(default=None) def __init__( self, @@ -138,7 +138,7 @@ def __init__( # during deployment deepcopy. It is initialized lazily in the property below. @property - def extension_client(self) -> "Extension": + def extension_client(self) -> 'Extension': """Lazy loads the Vertex AI Extension client.""" if self._code_interpreter_extension is None: self._code_interpreter_extension = _get_code_interpreter_extension( diff --git a/tests/unittests/code_executors/test_vertex_ai_code_executor.py b/tests/unittests/code_executors/test_vertex_ai_code_executor.py index 3d58ace28c..52f162b785 100644 --- a/tests/unittests/code_executors/test_vertex_ai_code_executor.py +++ b/tests/unittests/code_executors/test_vertex_ai_code_executor.py @@ -20,7 +20,9 @@ def setUp(self): ) self.executor = VertexAiCodeExecutor(resource_name=self.mock_resource_name) - def _create_mock_files(self, file_data: List[tuple[str, str, str]]) -> List[File]: + def _create_mock_files( + self, file_data: List[tuple[str, str, str]] + ) -> List[File]: """Helper to create File objects from (name, content, mime_type).""" return [ File(name=name, content=content, mime_type=mime_type)