fix(code_executors): correct sandbox JSON field names for stdout & input files#4516
fix(code_executors): correct sandbox JSON field names for stdout & input files#4516jbets4 wants to merge 1 commit intogoogle:mainfrom
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
Summary of ChangesHello @jbets4, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request addresses critical issues in the Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
|
Response from ADK Triaging Agent Hello @jbets4, thank you for your contribution! Before we can merge this pull request, you'll need to sign the Contributor License Agreement (CLA). You can find more information at https://cla.developers.google.com/. Thanks! |
There was a problem hiding this comment.
Code Review
This pull request correctly fixes a bug where incorrect JSON field names were used for sandbox code execution, affecting both input file uploads and output parsing. The changes to use content/mime_type for inputs and parse msg_out/msg_err for outputs (while maintaining backward compatibility) are well-implemented. The new unit tests are a great addition and validate the core fixes. However, I've noted a regression in test coverage due to the removal of a previous test case. My review includes a suggestion to add back tests for output file parsing and backward compatibility to ensure the code remains robust.
| @patch("vertexai.Client") | ||
| def test_execute_code_success( | ||
| self, | ||
| mock_vertexai_client, | ||
| mock_invocation_context, | ||
| def test_execute_code_parses_msg_out_msg_err( | ||
| self, | ||
| mock_vertexai_client, | ||
| mock_invocation_context, | ||
| ): | ||
| # Setup Mocks | ||
| # 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": "hello world", "stderr": ""} | ||
| {"msg_out": "hello", "msg_err": "boom"} | ||
| ).encode("utf-8") | ||
| mock_json_output.metadata = None | ||
| mock_response.outputs = [mock_json_output] | ||
|
|
||
| mock_file_output = MagicMock() | ||
| mock_file_output.mime_type = "text/plain" | ||
| mock_file_output.data = b"file content" | ||
| mock_file_output.metadata = MagicMock() | ||
| mock_file_output.metadata.attributes = {"file_name": b"file.txt"} | ||
|
|
||
| mock_png_file_output = MagicMock() | ||
| mock_png_file_output.mime_type = "image/png" | ||
| sample_png_bytes = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00\x00\x00\x1f\x15\xc4\x89" | ||
| mock_png_file_output.data = sample_png_bytes | ||
| mock_png_file_output.metadata = MagicMock() | ||
| mock_png_file_output.metadata.attributes = {"file_name": b"file.png"} | ||
|
|
||
| mock_response.outputs = [ | ||
| mock_json_output, | ||
| mock_file_output, | ||
| mock_png_file_output, | ||
| ] | ||
| mock_api_client.agent_engines.sandboxes.execute_code.return_value = ( | ||
| mock_response | ||
| mock_response | ||
| ) | ||
|
|
||
| # Execute | ||
| executor = AgentEngineSandboxCodeExecutor( | ||
| sandbox_resource_name="projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789" | ||
| sandbox_resource_name=( | ||
| "projects/123/locations/us-central1/" | ||
| "reasoningEngines/456/sandboxEnvironments/789" | ||
| ) | ||
| ) | ||
| result = executor.execute_code( | ||
| mock_invocation_context, | ||
| CodeExecutionInput(code='print("x")'), | ||
| ) | ||
| code_input = CodeExecutionInput(code='print("hello world")') | ||
| result = executor.execute_code(mock_invocation_context, code_input) | ||
|
|
||
| # Assert | ||
| assert result.stdout == "hello world" | ||
| assert not result.stderr | ||
| assert result.output_files[0].mime_type == "text/plain" | ||
| assert result.output_files[0].content == b"file content" | ||
|
|
||
| assert result.output_files[0].name == "file.txt" | ||
| assert result.output_files[1].mime_type == "image/png" | ||
| assert result.output_files[1].name == "file.png" | ||
| assert result.output_files[1].content == sample_png_bytes | ||
| mock_api_client.agent_engines.sandboxes.execute_code.assert_called_once_with( | ||
| name="projects/123/locations/us-central1/reasoningEngines/456/sandboxEnvironments/789", | ||
| input_data={"code": 'print("hello world")'}, | ||
| assert result.stdout == "hello" | ||
| assert result.stderr == "boom" | ||
|
|
||
|
|
||
| @patch("vertexai.Client") | ||
| def test_execute_code_input_files_use_correct_payload_keys( | ||
| self, | ||
| mock_vertexai_client, | ||
| mock_invocation_context, | ||
| ): | ||
| # 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": ""}).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" | ||
| ) | ||
| ) | ||
| code_input = CodeExecutionInput( | ||
| code='print("x")', | ||
| input_files=[ | ||
| File( | ||
| name="data.csv", | ||
| content=b"a,b\n1,2\n", | ||
| mime_type="text/csv", | ||
| ) | ||
| ], | ||
| ) | ||
|
|
||
| executor.execute_code(mock_invocation_context, code_input) | ||
|
|
||
| # Assert payload keys | ||
| call_kwargs = ( | ||
| mock_api_client.agent_engines.sandboxes.execute_code.call_args.kwargs | ||
| ) | ||
| files = call_kwargs["input_data"]["files"] | ||
| f0 = files[0] | ||
|
|
||
| assert f0["content"] == b"a,b\n1,2\n" | ||
| assert f0["mime_type"] == "text/csv" | ||
| assert "contents" not in f0 | ||
| assert "mimeType" not in f0 | ||
|
|
||
|
|
There was a problem hiding this comment.
While the new tests are excellent for verifying the primary fixes, the removal of the original test_execute_code_success has led to a loss of test coverage for a couple of important scenarios:
- Backward Compatibility: There is no longer a test to ensure that responses using the old
stdoutandstderrkeys are still parsed correctly. - Output File Parsing: The logic for parsing output files from the sandbox response is no longer covered by tests.
To prevent future regressions, it would be beneficial to add tests for these cases. Here are some suggestions:
@patch("vertexai.Client")
def test_execute_code_parses_stdout_stderr_for_backward_compatibility(
self,
mock_vertexai_client,
mock_invocation_context,
):
# 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": "old hello", "stderr": "old boom"}
).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"
)
)
result = executor.execute_code(
mock_invocation_context,
CodeExecutionInput(code='print("x")'),
)
# Assert
assert result.stdout == "old hello"
assert result.stderr == "old boom"
@patch("vertexai.Client")
def test_execute_code_parses_output_files(
self,
mock_vertexai_client,
mock_invocation_context,
):
# 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": ""}).encode("utf-8")
mock_json_output.metadata = None
mock_file_output = MagicMock()
mock_file_output.mime_type = "text/plain"
mock_file_output.data = b"file content"
mock_file_output.metadata = MagicMock()
mock_file_output.metadata.attributes = {"file_name": b"file.txt"}
mock_response.outputs = [mock_json_output, mock_file_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"
)
)
result = executor.execute_code(
mock_invocation_context,
CodeExecutionInput(code='print("x")'),
)
# Assert
assert len(result.output_files) == 1
assert result.output_files[0].name == "file.txt"
assert result.output_files[0].content == b"file content"
assert result.output_files[0].mime_type == "text/plain"
#3690
- Closes: #3690
Problem:
the AgentEngineSandboxCodeExecutor file was incorrectly using JSON field names when sending input files and when parsing the sandbox JSON response. This causes input files to be unreadable in the sandbox and caused stdout/stderr to appear empty.
Solution:
Testing Plan
Summary of results below
======================== 5 passed, 3 warnings in 4.28s =========================
PASSED [ 20%]PASSED [ 40%]PASSED [ 60%]PASSED [ 80%]PASSED [100%]
Process finished with exit code 0
Manual E2E test are beyond the scope for this change.
Checklist