fix(flows): reject malformed tool confirmation payloads fail-closed#4577
fix(flows): reject malformed tool confirmation payloads fail-closed#4577davidahmann wants to merge 1 commit intogoogle:mainfrom
Conversation
|
This PR hardens tool confirmation parsing to fail closed: malformed JSON or schema-invalid payloads now default to This contribution was informed by patterns from Gait: https://github.com/Clyra-AI/gait |
Summary of ChangesHello @davidahmann, 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 enhances the robustness of tool confirmation flows by introducing comprehensive error handling for malformed payloads. The change ensures that any invalid or unparseable confirmation data will result in a 'fail-closed' state, preventing ambiguous behavior and reinforcing explicit safety controls, aligning with the goal of creating replay/audit-friendly contracts. 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
|
There was a problem hiding this comment.
Code Review
The pull request effectively addresses the problem of malformed tool confirmation payloads by implementing robust error handling and ensuring a fail-closed mechanism. The addition of ValidationError and json.JSONDecodeError handling, along with a regression test, significantly improves the reliability and security of the confirmation parsing. The changes are well-structured and align with the stated objective.
| try: | ||
| tool_confirmation = ToolConfirmation.model_validate( | ||
| json.loads(function_response.response['response']) | ||
| ) | ||
| except ( | ||
| json.JSONDecodeError, | ||
| TypeError, | ||
| ValidationError, | ||
| ) as parse_err: | ||
| logger.warning( | ||
| 'Malformed tool confirmation payload for' | ||
| ' function_response_id=%s: %s', | ||
| function_response.id, | ||
| parse_err, | ||
| ) | ||
| tool_confirmation = ToolConfirmation(confirmed=False) |
There was a problem hiding this comment.
This try-except block correctly catches potential json.JSONDecodeError, TypeError, and ValidationError during the parsing of the tool confirmation payload. In case of an error, it logs a warning and defaults tool_confirmation to ToolConfirmation(confirmed=False), ensuring a fail-closed behavior as intended. This is a critical improvement for handling malformed inputs.
| from typing import TYPE_CHECKING | ||
|
|
||
| from google.genai import types | ||
| from pydantic import ValidationError |
| try: | ||
| tool_confirmation = ToolConfirmation.model_validate( | ||
| function_response.response | ||
| ) | ||
| except ValidationError as parse_err: | ||
| logger.warning( | ||
| 'Malformed tool confirmation payload for' | ||
| ' function_response_id=%s: %s', | ||
| function_response.id, | ||
| parse_err, | ||
| ) | ||
| tool_confirmation = ToolConfirmation(confirmed=False) |
There was a problem hiding this comment.
| @pytest.mark.asyncio | ||
| async def test_malformed_confirmation_payload_is_rejected_fail_closed( | ||
| self, | ||
| runner: testing_utils.InMemoryRunner, | ||
| agent: LlmAgent, | ||
| ): | ||
| """Malformed confirmation payloads must fail closed and reject tool calls.""" | ||
| initial_events = await runner.run_async(testing_utils.UserContent("test")) | ||
| ask_for_confirmation_function_call_id = ( | ||
| initial_events[1].content.parts[0].function_call.id | ||
| ) | ||
|
|
||
| malformed_confirmation = testing_utils.UserContent( | ||
| Part( | ||
| function_response=FunctionResponse( | ||
| id=ask_for_confirmation_function_call_id, | ||
| name=REQUEST_CONFIRMATION_FUNCTION_CALL_NAME, | ||
| response={"response": "{not-json"}, | ||
| ) | ||
| ) | ||
| ) | ||
|
|
||
| events = await runner.run_async(malformed_confirmation) | ||
| simplified = testing_utils.simplify_events(copy.deepcopy(events)) | ||
| assert simplified[0][1] == Part( | ||
| function_response=FunctionResponse( | ||
| name=agent.tools[0].name, | ||
| response={"error": "This tool call is rejected."}, | ||
| ) | ||
| ) |
Problem
Malformed tool confirmation payloads can bypass robust validation paths and risk ambiguous behavior at tool-execution decision points.
Why now
Gap item #4 focuses on explicit safety controls and replay/audit-friendly contracts. Confirmation parsing must fail closed whenever payload shape/content is invalid.
What changed
confirmed=Falseand emit warning log, ensuring rejection path is deterministic.Validation
uv run pyink --check --diff src/google/adk/flows/llm_flows/request_confirmation.py tests/unittests/runners/test_run_tool_confirmation.py✅uv run pytest tests/unittests/runners/test_run_tool_confirmation.py -k 'malformed_confirmation_payload_is_rejected_fail_closed or confirmation_flow' -q✅ (5 passed, 3 deselected)bash ./autoformat.sh✅ (repo-required format pass; introduced unrelated changes undercontributing/samples/gepa, intentionally excluded from this PR)Refs #4576