diff --git a/src/google/adk/code_executors/agent_engine_sandbox_code_executor.py b/src/google/adk/code_executors/agent_engine_sandbox_code_executor.py index f601d0455a..aa2f22ea73 100644 --- a/src/google/adk/code_executors/agent_engine_sandbox_code_executor.py +++ b/src/google/adk/code_executors/agent_engine_sandbox_code_executor.py @@ -28,7 +28,7 @@ from .code_execution_utils import CodeExecutionResult from .code_execution_utils import File -logger = logging.getLogger('google_adk.' + __name__) +logger = logging.getLogger("google_adk." + __name__) class AgentEngineSandboxCodeExecutor(BaseCodeExecutor): @@ -63,8 +63,8 @@ def __init__( **data: Additional keyword arguments to be passed to the base class. """ super().__init__(**data) - sandbox_resource_name_pattern = r'^projects/([a-zA-Z0-9-_]+)/locations/([a-zA-Z0-9-_]+)/reasoningEngines/(\d+)/sandboxEnvironments/(\d+)$' - agent_engine_resource_name_pattern = r'^projects/([a-zA-Z0-9-_]+)/locations/([a-zA-Z0-9-_]+)/reasoningEngines/(\d+)$' + sandbox_resource_name_pattern = r"^projects/([a-zA-Z0-9-_]+)/locations/([a-zA-Z0-9-_]+)/reasoningEngines/(\d+)/sandboxEnvironments/(\d+)$" + agent_engine_resource_name_pattern = r"^projects/([a-zA-Z0-9-_]+)/locations/([a-zA-Z0-9-_]+)/reasoningEngines/(\d+)$" if sandbox_resource_name is not None: self.sandbox_resource_name = sandbox_resource_name @@ -84,17 +84,17 @@ def __init__( # @TODO - Add TTL for sandbox creation after it is available # in SDK. operation = self._get_api_client().agent_engines.sandboxes.create( - spec={'code_execution_environment': {}}, + spec={"code_execution_environment": {}}, name=agent_engine_resource_name, config=types.CreateAgentEngineSandboxConfig( - display_name='default_sandbox' + display_name="default_sandbox" ), ) self.sandbox_resource_name = operation.response.name else: raise ValueError( - 'Either sandbox_resource_name or agent_engine_resource_name must be' - ' set.' + "Either sandbox_resource_name or agent_engine_resource_name must be" + " set." ) @override @@ -105,14 +105,14 @@ def execute_code( ) -> CodeExecutionResult: # Execute the code. input_data = { - 'code': code_execution_input.code, + "code": code_execution_input.code, } if code_execution_input.input_files: - input_data['files'] = [ + input_data["files"] = [ { - 'name': f.name, - 'contents': f.content, - 'mimeType': f.mime_type, + "name": f.name, + "content": f.content, + "mime_type": f.mime_type, } for f in code_execution_input.input_files ] @@ -123,27 +123,39 @@ def execute_code( input_data=input_data, ) ) - logger.debug('Executed code:\n```\n%s\n```', code_execution_input.code) + logger.debug("Executed code:\n```\n%s\n```", code_execution_input.code) saved_files = [] - stdout = '' - stderr = '' + stdout = "" + stderr = "" for output in code_execution_response.outputs: - if output.mime_type == 'application/json' and ( + if output.mime_type == "application/json" and ( output.metadata is None or output.metadata.attributes is None - or 'file_name' not in output.metadata.attributes + or "file_name" not in output.metadata.attributes ): - json_output_data = json.loads(output.data.decode('utf-8')) - stdout = json_output_data.get('stdout', '') - stderr = json_output_data.get('stderr', '') + json_output_data = json.loads(output.data.decode("utf-8")) + if isinstance(json_output_data, dict): + # Primary fields returned by the API are msg_out/msg_err. + # Fall back to stdout/stderr for backward compatibility. + stdout = json_output_data.get("msg_out") + if stdout is None: + stdout = json_output_data.get("stdout", "") + stderr = json_output_data.get("msg_err") + if stderr is None: + stderr = json_output_data.get("stderr", "") + else: + logger.warning( + "Received non-dict JSON output from sandbox: %s", + json_output_data, + ) else: - file_name = '' + file_name = "" if ( output.metadata is not None and output.metadata.attributes is not None ): - file_name = output.metadata.attributes.get('file_name', b'').decode( - 'utf-8' + file_name = output.metadata.attributes.get("file_name", b"").decode( + "utf-8" ) mime_type = output.mime_type if not mime_type: @@ -183,6 +195,6 @@ def _get_project_id_and_location_from_resource_name( match = re.fullmatch(pattern, resource_name) if not match: - raise ValueError(f'resource name {resource_name} is not valid.') + raise ValueError(f"resource name {resource_name} is not valid.") return match.groups()[0], match.groups()[1] diff --git a/tests/unittests/code_executors/test_agent_engine_sandbox_code_executor.py b/tests/unittests/code_executors/test_agent_engine_sandbox_code_executor.py index c948060184..66c40266bc 100644 --- a/tests/unittests/code_executors/test_agent_engine_sandbox_code_executor.py +++ b/tests/unittests/code_executors/test_agent_engine_sandbox_code_executor.py @@ -19,6 +19,7 @@ from google.adk.agents.invocation_context import InvocationContext from google.adk.code_executors.agent_engine_sandbox_code_executor import AgentEngineSandboxCodeExecutor from google.adk.code_executors.code_execution_utils import CodeExecutionInput +from google.adk.code_executors.code_execution_utils import File import pytest @@ -118,3 +119,414 @@ def test_execute_code_success( name="projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789", input_data={"code": 'print("hello world")'}, ) + + @patch("vertexai.Client") + def test_execute_code_with_msg_out_msg_err( + self, + mock_vertexai_client, + mock_invocation_context, + ): + """Tests that msg_out and msg_err fields from API response are parsed correctly.""" + # Setup Mocks + mock_api_client = MagicMock() + mock_vertexai_client.return_value = mock_api_client + mock_response = MagicMock() + mock_json_output = MagicMock() + mock_json_output.mime_type = "application/json" + mock_json_output.data = json.dumps( + {"msg_out": "hello from msg_out", "msg_err": "error from msg_err"} + ).encode("utf-8") + mock_json_output.metadata = None + + mock_response.outputs = [mock_json_output] + mock_api_client.agent_engines.sandboxes.execute_code.return_value = ( + mock_response + ) + + # Execute + executor = AgentEngineSandboxCodeExecutor( + sandbox_resource_name="projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789" + ) + code_input = CodeExecutionInput(code='print("hello")') + result = executor.execute_code(mock_invocation_context, code_input) + + # Assert - msg_out/msg_err should be mapped to stdout/stderr + assert result.stdout == "hello from msg_out" + assert result.stderr == "error from msg_err" + + @patch("vertexai.Client") + def test_execute_code_fallback_to_stdout_stderr( + self, + mock_vertexai_client, + mock_invocation_context, + ): + """Tests fallback to stdout/stderr when msg_out/msg_err are not present.""" + # Setup Mocks + mock_api_client = MagicMock() + mock_vertexai_client.return_value = mock_api_client + mock_response = MagicMock() + mock_json_output = MagicMock() + mock_json_output.mime_type = "application/json" + mock_json_output.data = json.dumps( + {"stdout": "fallback stdout", "stderr": "fallback stderr"} + ).encode("utf-8") + mock_json_output.metadata = None + + mock_response.outputs = [mock_json_output] + mock_api_client.agent_engines.sandboxes.execute_code.return_value = ( + mock_response + ) + + # Execute + executor = AgentEngineSandboxCodeExecutor( + sandbox_resource_name="projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789" + ) + code_input = CodeExecutionInput(code='print("hello")') + result = executor.execute_code(mock_invocation_context, code_input) + + # Assert - should fall back to stdout/stderr + assert result.stdout == "fallback stdout" + assert result.stderr == "fallback stderr" + + @patch("vertexai.Client") + def test_execute_code_msg_out_takes_precedence_over_stdout( + self, + mock_vertexai_client, + mock_invocation_context, + ): + """Tests that msg_out takes precedence over stdout when both are present.""" + # Setup Mocks + mock_api_client = MagicMock() + mock_vertexai_client.return_value = mock_api_client + mock_response = MagicMock() + mock_json_output = MagicMock() + mock_json_output.mime_type = "application/json" + # Both msg_out and stdout present - msg_out should win + mock_json_output.data = json.dumps({ + "msg_out": "primary output", + "msg_err": "primary error", + "stdout": "fallback output", + "stderr": "fallback error", + }).encode("utf-8") + mock_json_output.metadata = None + + mock_response.outputs = [mock_json_output] + mock_api_client.agent_engines.sandboxes.execute_code.return_value = ( + mock_response + ) + + # Execute + executor = AgentEngineSandboxCodeExecutor( + sandbox_resource_name="projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789" + ) + code_input = CodeExecutionInput(code='print("hello")') + result = executor.execute_code(mock_invocation_context, code_input) + + # Assert - msg_out/msg_err take precedence + assert result.stdout == "primary output" + assert result.stderr == "primary error" + + @patch("vertexai.Client") + def test_execute_code_partial_response_only_msg_out( + self, + mock_vertexai_client, + mock_invocation_context, + ): + """Tests handling when only msg_out is present (no msg_err).""" + # Setup Mocks + mock_api_client = MagicMock() + mock_vertexai_client.return_value = mock_api_client + mock_response = MagicMock() + mock_json_output = MagicMock() + mock_json_output.mime_type = "application/json" + mock_json_output.data = json.dumps({"msg_out": "only output"}).encode( + "utf-8" + ) + mock_json_output.metadata = None + + mock_response.outputs = [mock_json_output] + mock_api_client.agent_engines.sandboxes.execute_code.return_value = ( + mock_response + ) + + # Execute + executor = AgentEngineSandboxCodeExecutor( + sandbox_resource_name="projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789" + ) + code_input = CodeExecutionInput(code='print("hello")') + result = executor.execute_code(mock_invocation_context, code_input) + + # Assert - stdout populated, stderr empty string (not None) + assert result.stdout == "only output" + assert result.stderr == "" + + @patch("vertexai.Client") + def test_execute_code_partial_response_only_msg_err( + self, + mock_vertexai_client, + mock_invocation_context, + ): + """Tests handling when only msg_err is present (no msg_out).""" + # Setup Mocks + mock_api_client = MagicMock() + mock_vertexai_client.return_value = mock_api_client + mock_response = MagicMock() + mock_json_output = MagicMock() + mock_json_output.mime_type = "application/json" + mock_json_output.data = json.dumps({"msg_err": "only error"}).encode( + "utf-8" + ) + mock_json_output.metadata = None + + mock_response.outputs = [mock_json_output] + mock_api_client.agent_engines.sandboxes.execute_code.return_value = ( + mock_response + ) + + # Execute + executor = AgentEngineSandboxCodeExecutor( + sandbox_resource_name="projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789" + ) + code_input = CodeExecutionInput(code='print("hello")') + result = executor.execute_code(mock_invocation_context, code_input) + + # Assert - stderr populated, stdout empty string (not None) + assert result.stdout == "" + assert result.stderr == "only error" + + @patch("vertexai.Client") + def test_execute_code_empty_response( + self, + mock_vertexai_client, + mock_invocation_context, + ): + """Tests handling when response has no stdout/stderr fields.""" + # Setup Mocks + mock_api_client = MagicMock() + mock_vertexai_client.return_value = mock_api_client + mock_response = MagicMock() + mock_json_output = MagicMock() + mock_json_output.mime_type = "application/json" + mock_json_output.data = json.dumps({}).encode("utf-8") + mock_json_output.metadata = None + + mock_response.outputs = [mock_json_output] + mock_api_client.agent_engines.sandboxes.execute_code.return_value = ( + mock_response + ) + + # Execute + executor = AgentEngineSandboxCodeExecutor( + sandbox_resource_name="projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789" + ) + code_input = CodeExecutionInput(code='print("hello")') + result = executor.execute_code(mock_invocation_context, code_input) + + # Assert - both should be empty strings + assert result.stdout == "" + assert result.stderr == "" + + @patch("vertexai.Client") + def test_execute_code_with_input_files_uses_correct_payload_keys( + self, + mock_vertexai_client, + mock_invocation_context, + ): + """Tests that input files use content and mime_type keys (not contents/mimeType).""" + # Setup Mocks + mock_api_client = MagicMock() + mock_vertexai_client.return_value = mock_api_client + mock_response = MagicMock() + mock_json_output = MagicMock() + mock_json_output.mime_type = "application/json" + mock_json_output.data = json.dumps({"msg_out": "ok", "msg_err": ""}).encode( + "utf-8" + ) + mock_json_output.metadata = None + mock_response.outputs = [mock_json_output] + mock_api_client.agent_engines.sandboxes.execute_code.return_value = ( + mock_response + ) + + # Execute with input files + executor = AgentEngineSandboxCodeExecutor( + sandbox_resource_name="projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789" + ) + input_files = [ + File(name="data.csv", content=b"col1,col2\n1,2", mime_type="text/csv"), + File( + name="config.json", + content=b'{"key": "value"}', + mime_type="application/json", + ), + ] + code_input = CodeExecutionInput( + code='import pandas as pd; df = pd.read_csv("data.csv")', + input_files=input_files, + ) + executor.execute_code(mock_invocation_context, code_input) + + # Assert - verify the payload uses correct keys + call_args = mock_api_client.agent_engines.sandboxes.execute_code.call_args + input_data = call_args.kwargs["input_data"] + + assert "files" in input_data + assert len(input_data["files"]) == 2 + + # Check first file + file1 = input_data["files"][0] + assert file1["name"] == "data.csv" + assert file1["content"] == b"col1,col2\n1,2" + assert file1["mime_type"] == "text/csv" + # Ensure old keys are NOT present + assert "contents" not in file1 + assert "mimeType" not in file1 + + # Check second file + file2 = input_data["files"][1] + assert file2["name"] == "config.json" + assert file2["content"] == b'{"key": "value"}' + assert file2["mime_type"] == "application/json" + assert "contents" not in file2 + assert "mimeType" not in file2 + + @patch("vertexai.Client") + def test_execute_code_without_input_files_no_files_key( + self, + mock_vertexai_client, + mock_invocation_context, + ): + """Tests that files key is not present when no input files are provided.""" + # Setup Mocks + mock_api_client = MagicMock() + mock_vertexai_client.return_value = mock_api_client + mock_response = MagicMock() + mock_json_output = MagicMock() + mock_json_output.mime_type = "application/json" + mock_json_output.data = json.dumps({"msg_out": "ok"}).encode("utf-8") + mock_json_output.metadata = None + mock_response.outputs = [mock_json_output] + mock_api_client.agent_engines.sandboxes.execute_code.return_value = ( + mock_response + ) + + # Execute without input files + executor = AgentEngineSandboxCodeExecutor( + sandbox_resource_name="projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789" + ) + code_input = CodeExecutionInput(code='print("hello")') + executor.execute_code(mock_invocation_context, code_input) + + # Assert - verify no files key in payload + call_args = mock_api_client.agent_engines.sandboxes.execute_code.call_args + input_data = call_args.kwargs["input_data"] + assert "files" not in input_data + assert input_data == {"code": 'print("hello")'} + + @patch("vertexai.Client") + def test_execute_code_output_files_metadata_preserved( + self, + mock_vertexai_client, + mock_invocation_context, + ): + """Tests that output file metadata is correctly preserved.""" + # Setup Mocks + mock_api_client = MagicMock() + mock_vertexai_client.return_value = mock_api_client + mock_response = MagicMock() + + mock_json_output = MagicMock() + mock_json_output.mime_type = "application/json" + mock_json_output.data = json.dumps({"msg_out": "done"}).encode("utf-8") + mock_json_output.metadata = None + + mock_csv_output = MagicMock() + mock_csv_output.mime_type = "text/csv" + mock_csv_output.data = b"a,b,c\n1,2,3" + mock_csv_output.metadata = MagicMock() + mock_csv_output.metadata.attributes = {"file_name": b"output.csv"} + + mock_response.outputs = [mock_json_output, mock_csv_output] + mock_api_client.agent_engines.sandboxes.execute_code.return_value = ( + mock_response + ) + + # Execute + executor = AgentEngineSandboxCodeExecutor( + sandbox_resource_name="projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789" + ) + code_input = CodeExecutionInput(code='df.to_csv("output.csv")') + result = executor.execute_code(mock_invocation_context, code_input) + + # Assert + assert result.stdout == "done" + assert len(result.output_files) == 1 + assert result.output_files[0].name == "output.csv" + assert result.output_files[0].content == b"a,b,c\n1,2,3" + assert result.output_files[0].mime_type == "text/csv" + + @patch("vertexai.Client") + def test_execute_code_non_dict_json_response_logs_warning( + self, + mock_vertexai_client, + mock_invocation_context, + ): + """Tests that non-dict JSON responses are handled gracefully with a warning.""" + # Setup Mocks + mock_api_client = MagicMock() + mock_vertexai_client.return_value = mock_api_client + mock_response = MagicMock() + mock_json_output = MagicMock() + mock_json_output.mime_type = "application/json" + # Return a list instead of a dict - this is valid JSON but not expected + mock_json_output.data = json.dumps(["unexpected", "list"]).encode("utf-8") + mock_json_output.metadata = None + + mock_response.outputs = [mock_json_output] + mock_api_client.agent_engines.sandboxes.execute_code.return_value = ( + mock_response + ) + + # Execute + executor = AgentEngineSandboxCodeExecutor( + sandbox_resource_name="projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789" + ) + code_input = CodeExecutionInput(code='print("hello")') + result = executor.execute_code(mock_invocation_context, code_input) + + # Assert - should handle gracefully with empty stdout/stderr + assert result.stdout == "" + assert result.stderr == "" + + @patch("vertexai.Client") + def test_execute_code_string_json_response_logs_warning( + self, + mock_vertexai_client, + mock_invocation_context, + ): + """Tests that string JSON responses are handled gracefully.""" + # Setup Mocks + mock_api_client = MagicMock() + mock_vertexai_client.return_value = mock_api_client + mock_response = MagicMock() + mock_json_output = MagicMock() + mock_json_output.mime_type = "application/json" + # Return a string instead of a dict + mock_json_output.data = json.dumps("just a string").encode("utf-8") + mock_json_output.metadata = None + + mock_response.outputs = [mock_json_output] + mock_api_client.agent_engines.sandboxes.execute_code.return_value = ( + mock_response + ) + + # Execute + executor = AgentEngineSandboxCodeExecutor( + sandbox_resource_name="projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789" + ) + code_input = CodeExecutionInput(code='print("hello")') + result = executor.execute_code(mock_invocation_context, code_input) + + # Assert - should handle gracefully with empty stdout/stderr + assert result.stdout == "" + assert result.stderr == ""