From 9715a7503abe9e8c9d999c7783c34050e303ec63 Mon Sep 17 00:00:00 2001 From: Lucas Date: Wed, 2 Jul 2025 12:27:46 -0700 Subject: [PATCH] fix(code-executor): resolve deepcopy recursion in VertexAiCodeExecutor Add __deepcopy__ method to handle extension serialization during agent deployment. Extension gets re-initialized automatically when needed. Fixes agent engine deployment failures. --- .../code_executors/vertex_ai_code_executor.py | 35 +++ .../test_vertex_ai_code_executor.py | 213 ++++++++++++++++++ 2 files changed, 248 insertions(+) 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 b1dd58ce7..6068bfa77 100644 --- a/src/google/adk/code_executors/vertex_ai_code_executor.py +++ b/src/google/adk/code_executors/vertex_ai_code_executor.py @@ -136,10 +136,42 @@ def __init__( """ super().__init__(**data) self.resource_name = resource_name + self.initialize_extension() + + def initialize_extension(self) -> None: + """Initializes the Vertex Code Interpreter Extension.""" self._code_interpreter_extension = _get_code_interpreter_extension( self.resource_name ) + def _ensure_extension_initialized(self) -> None: + """Ensures the extension is initialized, re-initializing if necessary.""" + if ( + not hasattr(self, '_code_interpreter_extension') + or self._code_interpreter_extension is None + ): + self.initialize_extension() + + def __deepcopy__(self, memo: dict[int, Any] | None = None): + # Create a copy by temporarily removing the problematic extension + # Store the extension temporarily + original_extension = getattr(self, '_code_interpreter_extension', None) + + # Temporarily set extension to None + self._code_interpreter_extension = None + + # Now perform the deepcopy safely + try: + copied = super().__deepcopy__(memo) + finally: + # Restore the original extension + self._code_interpreter_extension = original_extension + + # Set the copied object's extension to None - it will be re-initialized when needed + copied._code_interpreter_extension = None + + return copied + @override def execute_code( self, @@ -209,6 +241,9 @@ def _execute_code_interpreter( Returns: The response from the code interpreter extension. """ + # Ensure extension is initialized (re-initialize if it was set to None during deepcopy) + self._ensure_extension_initialized() + operation_params = {'code': code} if input_files: operation_params['files'] = [ 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 000000000..3c9a34b02 --- /dev/null +++ b/tests/unittests/code_executors/test_vertex_ai_code_executor.py @@ -0,0 +1,213 @@ +""" +Test suite for VertexAiCodeExecutor deep copy and extension initialization fixes. + +This test validates the critical fixes made to address: +1. Deep copy recursion errors during agent engine deployment +2. Extension state management during serialization +3. Automatic extension re-initialization after deep copy +""" + +import copy +from unittest.mock import MagicMock +from unittest.mock import Mock +from unittest.mock import patch + +from google.adk.agents.invocation_context import InvocationContext +from google.adk.code_executors import VertexAiCodeExecutor +from google.adk.code_executors.code_execution_utils import CodeExecutionInput +from google.adk.code_executors.code_execution_utils import File +import pytest + + +class TestVertexAiCodeExecutorFixes: + """Test class for VertexAiCodeExecutor deep copy and extension fixes.""" + + @pytest.fixture + def mock_extension(self): + """Create a mock extension for testing.""" + mock_ext = Mock() + mock_ext.execute.return_value = { + "execution_result": 'print("Hello World")\nHello World', + "execution_error": "", + "output_files": [], + } + return mock_ext + + @pytest.fixture + def code_executor(self, mock_extension): + """Create a VertexAiCodeExecutor instance with mocked extension.""" + with patch( + "google.adk.code_executors.vertex_ai_code_executor._get_code_interpreter_extension", + return_value=mock_extension, + ): + executor = VertexAiCodeExecutor() + return executor + + def test_deep_copy_no_recursion_error(self, code_executor): + """Test that deep copy works without recursion errors.""" + # This was the core issue: deep copy would cause recursion errors + # due to the extension object not being serializable + try: + copied_executor = copy.deepcopy(code_executor) + assert copied_executor is not None + assert copied_executor != code_executor # Different instances + print("Deep copy completed without recursion errors") + except RecursionError: + pytest.fail("Deep copy still causes recursion error - fix not working") + + def test_extension_state_after_deep_copy(self, code_executor, mock_extension): + """Test that extension is properly managed after deep copy.""" + # Original executor should have extension + assert hasattr(code_executor, "_code_interpreter_extension") + + # After deep copy, the copied object should have None extension initially + copied_executor = copy.deepcopy(code_executor) + + # The copied executor's extension should be None (to be re-initialized) + assert copied_executor._code_interpreter_extension is None + + # Original executor should still have its extension + assert code_executor._code_interpreter_extension == mock_extension + print("Extension state properly managed during deep copy") + + def test_extension_re_initialization(self, mock_extension): + """Test that extension gets re-initialized when needed.""" + with patch( + "google.adk.code_executors.vertex_ai_code_executor._get_code_interpreter_extension", + return_value=mock_extension, + ) as mock_get_ext: + + executor = VertexAiCodeExecutor() + + # Deep copy the executor + copied_executor = copy.deepcopy(executor) + + # Extension should be None after deep copy + assert copied_executor._code_interpreter_extension is None + + # Call ensure_extension_initialized - should trigger re-initialization + copied_executor._ensure_extension_initialized() + + # Extension should now be re-initialized + assert copied_executor._code_interpreter_extension == mock_extension + print("Extension re-initialization working correctly") + + def test_code_execution_after_deep_copy(self, code_executor, mock_extension): + """Test that code execution works after deep copy.""" + # Create test input + invocation_context = Mock(spec=InvocationContext) + code_input = CodeExecutionInput( + code="print('Hello from copied executor')", execution_id="test-123" + ) + + # Deep copy the executor + copied_executor = copy.deepcopy(code_executor) + + # Execute code with copied executor - should trigger re-initialization + with patch( + "google.adk.code_executors.vertex_ai_code_executor._get_code_interpreter_extension", + return_value=mock_extension, + ): + result = copied_executor.execute_code(invocation_context, code_input) + + # Verify execution worked + assert result is not None + assert "Hello World" in result.stdout + + # Verify extension was re-initialized during execution + assert copied_executor._code_interpreter_extension == mock_extension + print("Code execution works after deep copy with auto re-initialization") + + def test_ensure_extension_initialized_idempotent( + self, code_executor, mock_extension + ): + """Test that _ensure_extension_initialized is safe to call multiple times.""" + original_extension = code_executor._code_interpreter_extension + + # Call multiple times + code_executor._ensure_extension_initialized() + code_executor._ensure_extension_initialized() + code_executor._ensure_extension_initialized() + + # Extension should remain the same + assert code_executor._code_interpreter_extension == original_extension + print("Extension initialization is idempotent") + + def test_extension_initialization_with_resource_name(self, mock_extension): + """Test extension initialization with custom resource name.""" + resource_name = "projects/test/locations/us-central1/extensions/123" + + with patch( + "google.adk.code_executors.vertex_ai_code_executor._get_code_interpreter_extension", + return_value=mock_extension, + ) as mock_get_ext: + + executor = VertexAiCodeExecutor(resource_name=resource_name) + + # Verify resource name was passed correctly + mock_get_ext.assert_called_with(resource_name) + assert executor.resource_name == resource_name + print("Resource name properly handled during initialization") + + @patch.dict( + "os.environ", {"CODE_INTERPRETER_EXTENSION_NAME": "test-extension"} + ) + def test_environment_variable_handling(self, mock_extension): + """Test that environment variables are properly handled.""" + with patch( + "google.adk.code_executors.vertex_ai_code_executor._get_code_interpreter_extension", + return_value=mock_extension, + ) as mock_get_ext: + + executor = VertexAiCodeExecutor() + + # Should use environment variable + mock_get_ext.assert_called_with(None) # No resource_name passed + print("Environment variable handling works correctly") + + +def test_integration_with_agent_engine_deployment(): + """Integration test simulating agent engine deployment process.""" + print("\nRunning integration test for agent engine deployment...") + + with patch( + "google.adk.code_executors.vertex_ai_code_executor._get_code_interpreter_extension" + ) as mock_get_ext: + mock_extension = Mock() + mock_extension.execute.return_value = { + "execution_result": "Deployment test successful", + "execution_error": "", + "output_files": [], + } + mock_get_ext.return_value = mock_extension + + # Create executor (simulating agent creation) + executor = VertexAiCodeExecutor() + + # Simulate agent engine deployment (which involves deep copying) + try: + serialized_executor = copy.deepcopy(executor) + + # Simulate code execution on deployed agent + invocation_context = Mock(spec=InvocationContext) + code_input = CodeExecutionInput( + code="print('Agent deployed successfully!')", + execution_id="deployment-test", + ) + + result = serialized_executor.execute_code(invocation_context, code_input) + + assert result is not None + assert "Deployment test successful" in result.stdout + print( + "Integration test passed - agent engine deployment simulation" + " successful" + ) + + except Exception as e: + pytest.fail(f"Integration test failed: {e}") + + +if __name__ == "__main__": + # Run the tests + pytest.main([__file__, "-v", "--tb=short"])