-
Notifications
You must be signed in to change notification settings - Fork 3k
fix(flows): reject malformed tool confirmation payloads fail-closed #4577
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -19,6 +19,7 @@ | |
| from typing import TYPE_CHECKING | ||
|
|
||
| from google.genai import types | ||
| from pydantic import ValidationError | ||
| from typing_extensions import override | ||
|
|
||
| from . import functions | ||
|
|
@@ -81,13 +82,35 @@ async def run_async( | |
| # ADK client must send a resuming run request with a function response | ||
| # that always encapsulate the confirmation result with a 'response' | ||
| # key | ||
| tool_confirmation = ToolConfirmation.model_validate( | ||
| json.loads(function_response.response['response']) | ||
| ) | ||
| 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) | ||
|
Comment on lines
+85
to
+100
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This |
||
| else: | ||
| tool_confirmation = ToolConfirmation.model_validate( | ||
| function_response.response | ||
| ) | ||
| 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) | ||
|
Comment on lines
+102
to
+113
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. |
||
| request_confirmation_function_responses[function_response.id] = ( | ||
| tool_confirmation | ||
| ) | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -224,6 +224,37 @@ async def test_confirmation_flow( | |
| == expected_parts_final | ||
| ) | ||
|
|
||
| @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."}, | ||
| ) | ||
| ) | ||
|
Comment on lines
+227
to
+256
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. |
||
|
|
||
|
|
||
| class TestHITLConfirmationFlowWithCustomPayloadSchema(BaseHITLTest): | ||
| """Tests the HITL confirmation flow with a single agent, for custom confirmation payload schema.""" | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The
ValidationErrorimport is correctly added to handle Pydantic validation errors, which is crucial for robust parsing of tool confirmation payloads.