chore: import upstream snapshot with attribution

tests/unittests/plugins/test_reflect_retry_tool_plugin.py

@@ -0,0 +1,668 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+from unittest import IsolatedAsyncioTestCase
+from unittest.mock import Mock
+
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.plugins.reflect_retry_tool_plugin import REFLECT_AND_RETRY_RESPONSE_TYPE
+from google.adk.plugins.reflect_retry_tool_plugin import ReflectAndRetryToolPlugin
+from google.adk.tools.base_tool import BaseTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+
+from .. import testing_utils
+
+
+class MockTool(BaseTool):
+  """Mock tool for testing purposes."""
+
+  def __init__(self, name: str = "mock_tool"):
+    self.name = name
+    self.description = f"Mock tool named {name}"
+
+  async def run(self, **kwargs) -> Any:
+    return "mock result"
+
+
+class CustomErrorExtractionPlugin(ReflectAndRetryToolPlugin):
+  """Custom plugin for testing error extraction from tool responses."""
+
+  def __init__(self, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+    self.error_conditions = {}
+
+  def set_error_condition(self, condition_func):
+    """Set a custom error condition function for testing."""
+    self.error_condition = condition_func
+
+  async def extract_error_from_result(
+      self, *, tool, tool_args, tool_context, result
+  ):
+    """Extract error based on custom conditions set for testing."""
+    if hasattr(self, "error_condition"):
+      return self.error_condition(result)
+    return None
+
+
+# Inheriting from IsolatedAsyncioTestCase ensures consistent behavior.
+# See https://github.com/pytest-dev/pytest-asyncio/issues/1039
+class TestReflectAndRetryToolPlugin(IsolatedAsyncioTestCase):
+  """Comprehensive tests for ReflectAndRetryToolPlugin focusing on behavior."""
+
+  def get_plugin(self):
+    """Create a default plugin instance for testing."""
+    return ReflectAndRetryToolPlugin()
+
+  def get_custom_plugin(self):
+    """Create a plugin with custom parameters."""
+    return ReflectAndRetryToolPlugin(
+        name="custom_plugin",
+        max_retries=5,
+        throw_exception_if_retry_exceeded=False,
+    )
+
+  def get_mock_tool(self):
+    """Create a mock tool for testing."""
+    return MockTool("test_tool_id")
+
+  def get_mock_tool_context(self):
+    """Create a mock tool context."""
+    return Mock(spec=ToolContext)
+
+  def get_custom_error_plugin(self):
+    """Create a custom error extraction plugin for testing."""
+    return CustomErrorExtractionPlugin(max_retries=3)
+
+  def get_sample_tool_args(self):
+    """Sample tool arguments for testing."""
+    return {"param1": "value1", "param2": 42, "param3": True}
+
+  async def test_plugin_initialization_default(self):
+    """Test plugin initialization with default parameters."""
+    plugin = self.get_plugin()
+
+    self.assertEqual(plugin.name, "reflect_retry_tool_plugin")
+    self.assertEqual(plugin.max_retries, 3)
+    self.assertIs(plugin.throw_exception_if_retry_exceeded, True)
+
+  async def test_plugin_initialization_custom(self):
+    """Test plugin initialization with custom parameters."""
+    plugin = ReflectAndRetryToolPlugin(
+        name="custom_name",
+        max_retries=10,
+        throw_exception_if_retry_exceeded=False,
+    )
+
+    self.assertEqual(plugin.name, "custom_name")
+    self.assertEqual(plugin.max_retries, 10)
+    self.assertIsNot(plugin.throw_exception_if_retry_exceeded, True)
+
+  async def test_after_tool_callback_successful_call(self):
+    """Test after_tool_callback with successful tool call."""
+    plugin = self.get_plugin()
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    result = {"success": True, "data": "test_data"}
+
+    callback_result = await plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result=result,
+    )
+
+    # Should return None for successful calls
+    self.assertIsNone(callback_result)
+
+  async def test_after_tool_callback_ignore_retry_response(self):
+    """Test that retry responses are ignored in after_tool_callback."""
+    plugin = self.get_plugin()
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    retry_result = {"response_type": REFLECT_AND_RETRY_RESPONSE_TYPE}
+
+    callback_result = await plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result=retry_result,
+    )
+
+    # Retry responses should be ignored
+    self.assertIsNone(callback_result)
+
+  async def test_on_tool_error_callback_max_retries_zero(self):
+    """Test error callback when max_retries is 0.
+
+    This should return None so that the exception is rethrown
+    """
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    plugin = ReflectAndRetryToolPlugin(max_retries=0)
+    error = ValueError("Test error")
+
+    with self.assertRaises(ValueError) as cm:
+      await plugin.on_tool_error_callback(
+          tool=mock_tool,
+          tool_args=sample_tool_args,
+          tool_context=mock_tool_context,
+          error=error,
+      )
+
+    # Should re-raise the original exception when max_retries is 0
+    self.assertIs(cm.exception, error)
+
+  async def test_on_tool_error_callback_max_retries_zero_without_exception(
+      self,
+  ):
+    """Test error callback when max_retries is 0 and exception is disabled."""
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    plugin = ReflectAndRetryToolPlugin(
+        max_retries=0, throw_exception_if_retry_exceeded=False
+    )
+    error = ValueError("Test error")
+
+    result = await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+
+    # Should return a retry exceeded message instead of raising
+    self.assertIsNotNone(result)
+    self.assertEqual(result["response_type"], REFLECT_AND_RETRY_RESPONSE_TYPE)
+    self.assertEqual(result["error_type"], "ValueError")
+    self.assertEqual(result["retry_count"], 0)
+    self.assertIn(
+        "the retry limit has been exceeded", result["reflection_guidance"]
+    )
+
+  async def test_on_tool_error_callback_max_retries_zero_with_dict_error(self):
+    """Test error callback when max_retries is 0 and error is a dict."""
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    plugin = CustomErrorExtractionPlugin(
+        max_retries=0, throw_exception_if_retry_exceeded=True
+    )
+    dict_error = {"status": "error", "message": "Custom dict error"}
+    plugin.set_error_condition(lambda result: dict_error)
+
+    with self.assertRaises(Exception) as cm:
+      await plugin.after_tool_callback(
+          tool=mock_tool,
+          tool_args=sample_tool_args,
+          tool_context=mock_tool_context,
+          result={"some": "result"},
+      )
+
+    # Should raise an Exception wrapping the dict
+    self.assertNotIsInstance(cm.exception, TypeError)
+    self.assertIn("Custom dict error", str(cm.exception))
+
+  async def test_on_tool_error_callback_first_failure(self):
+    """Test first tool failure creates reflection response."""
+    plugin = self.get_plugin()
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    error = ValueError("Test error message")
+
+    result = await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+
+    self.assertIsNotNone(result)
+    self.assertEqual(result["response_type"], REFLECT_AND_RETRY_RESPONSE_TYPE)
+    self.assertEqual(result["error_type"], "ValueError")
+    self.assertEqual(result["error_details"], "Test error message")
+    self.assertEqual(result["retry_count"], 1)
+    self.assertIn("test_tool_id", result["reflection_guidance"])
+    self.assertIn("Test error message", result["reflection_guidance"])
+
+  async def test_retry_behavior_with_consecutive_failures(self):
+    """Test the retry behavior with consecutive failures."""
+    plugin = self.get_plugin()
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    error = RuntimeError("Runtime error")
+
+    # First failure
+    result1 = await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+    self.assertEqual(result1["retry_count"], 1)
+
+    # Second failure - should have different retry count based on plugin logic
+    result2 = await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+    # The plugin's internal logic determines the exact retry count
+    self.assertIsNotNone(result2)
+    self.assertEqual(result2["response_type"], REFLECT_AND_RETRY_RESPONSE_TYPE)
+    self.assertEqual(result2["retry_count"], 2)
+
+  async def test_different_tools_behavior(self):
+    """Test behavior when using different tools."""
+    plugin = self.get_plugin()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    tool1 = MockTool("tool1")
+    tool2 = MockTool("tool2")
+    error = ValueError("Test error")
+
+    # First failure on tool1
+    result1 = await plugin.on_tool_error_callback(
+        tool=tool1,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+    self.assertEqual(result1["retry_count"], 1)
+
+    # Failure on tool2
+    result2 = await plugin.on_tool_error_callback(
+        tool=tool2,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+    # Since tool is different, retry count should start over.
+    self.assertIsNotNone(result2)
+    self.assertEqual(result2["response_type"], REFLECT_AND_RETRY_RESPONSE_TYPE)
+    self.assertEqual(result2["retry_count"], 1)
+
+  async def test_max_retries_exceeded_with_exception(self):
+    """Test that original exception is raised when max retries exceeded."""
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    plugin = ReflectAndRetryToolPlugin(
+        max_retries=1, throw_exception_if_retry_exceeded=True
+    )
+    error = ConnectionError("Connection failed")
+
+    # First call should succeed and return a retry response
+    await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+
+    # Second call should exceed max_retries and raise
+    with self.assertRaises(ConnectionError) as cm:
+      await plugin.on_tool_error_callback(
+          tool=mock_tool,
+          tool_args=sample_tool_args,
+          tool_context=mock_tool_context,
+          error=error,
+      )
+
+    # Verify exception properties
+    self.assertIs(cm.exception, error)
+
+  async def test_max_retries_exceeded_with_dict_error(self):
+    """Test that Exception is raised when max retries exceeded with dict error."""
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    plugin = CustomErrorExtractionPlugin(
+        max_retries=1, throw_exception_if_retry_exceeded=True
+    )
+    dict_error = {"status": "error", "message": "Custom dict error"}
+    plugin.set_error_condition(lambda result: dict_error)
+
+    # First call should fail and return a retry response
+    result1 = await plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result={"some": "result"},
+    )
+    self.assertIsNotNone(result1)
+    self.assertEqual(result1["retry_count"], 1)
+
+    # Second call should exceed max_retries and raise
+    with self.assertRaises(Exception) as cm:
+      await plugin.after_tool_callback(
+          tool=mock_tool,
+          tool_args=sample_tool_args,
+          tool_context=mock_tool_context,
+          result={"some": "result"},
+      )
+
+    # Verify exception properties
+    self.assertNotIsInstance(cm.exception, TypeError)
+    self.assertIn("Custom dict error", str(cm.exception))
+
+  async def test_max_retries_exceeded_without_exception(self):
+    """Test max retries exceeded returns failure message when exception is disabled."""
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    plugin = ReflectAndRetryToolPlugin(
+        max_retries=2, throw_exception_if_retry_exceeded=False
+    )
+    error = TimeoutError("Timeout occurred")
+
+    # Call until we exceed the retry limit
+    result = None
+    for _ in range(3):
+      result = await plugin.on_tool_error_callback(
+          tool=mock_tool,
+          tool_args=sample_tool_args,
+          tool_context=mock_tool_context,
+          error=error,
+      )
+
+    # Should get a retry exceeded message on the last call
+    self.assertIsNotNone(result)
+    self.assertEqual(result["response_type"], REFLECT_AND_RETRY_RESPONSE_TYPE)
+    self.assertEqual(result["error_type"], "TimeoutError")
+    self.assertIn(
+        "the retry limit has been exceeded", result["reflection_guidance"]
+    )
+    self.assertIn("Do not attempt to use the", result["reflection_guidance"])
+
+  async def test_successful_call_resets_retry_behavior(self):
+    """Test that successful calls reset the retry behavior."""
+    plugin = self.get_plugin()
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    error = ValueError("Test error")
+
+    # First failure
+    result1 = await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+    self.assertEqual(result1["retry_count"], 1)
+
+    # Successful call
+    await plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result={"success": True},
+    )
+
+    # Next failure should start fresh
+    result2 = await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+    self.assertEqual(result2["retry_count"], 1)  # Should restart from 1
+
+  async def test_none_result_handling(self):
+    """Test handling of None results in after_tool_callback."""
+    plugin = self.get_plugin()
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+
+    # None result should be handled gracefully
+    callback_result = await plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result=None,
+    )
+
+    self.assertIsNone(callback_result)
+
+  async def test_empty_tool_args_handling(self):
+    """Test handling of empty tool arguments."""
+    plugin = self.get_plugin()
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    empty_args = {}
+    error = ValueError("Test error")
+
+    result = await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=empty_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+
+    self.assertIsNotNone(result)
+    # Empty args should be represented in the response
+    self.assertIn("{}", result["reflection_guidance"])
+
+  async def test_retry_count_progression(self):
+    """Test that retry counts progress correctly for the same tool."""
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    plugin = ReflectAndRetryToolPlugin(max_retries=5)
+    error = ValueError("Test error")
+    tool = MockTool("single_tool")
+
+    for i in range(1, 4):
+      result = await plugin.on_tool_error_callback(
+          tool=tool,
+          tool_args=sample_tool_args,
+          tool_context=mock_tool_context,
+          error=error,
+      )
+      self.assertEqual(result["retry_count"], i)
+
+  async def test_max_retries_parameter_behavior(self):
+    """Test that max_retries parameter affects behavior correctly."""
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    # Test with very low max_retries
+    plugin = ReflectAndRetryToolPlugin(
+        max_retries=1, throw_exception_if_retry_exceeded=False
+    )
+    error = ValueError("Test error")
+
+    # First call is fine
+    await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+
+    # Second call exceeds limit
+    result = await plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=error,
+    )
+
+    # Should hit max retries quickly with max_retries=1
+    self.assertIn(
+        "the retry limit has been exceeded.", result["reflection_guidance"]
+    )
+
+  async def test_default_extract_error_returns_none(self):
+    """Test that default extract_error_from_result returns None."""
+    plugin = self.get_plugin()
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    result = {"status": "success", "data": "some data"}
+
+    error = await plugin.extract_error_from_result(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result=result,
+    )
+    self.assertIsNone(error)
+
+  async def test_custom_error_detection_and_success_handling(self):
+    """Test custom error detection, success handling, and retry progression."""
+    custom_error_plugin = self.get_custom_error_plugin()
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    custom_error_plugin.set_error_condition(
+        lambda result: result if result.get("status") == "error" else None
+    )
+
+    # Test error detection
+    error_result = {"status": "error", "message": "Something went wrong"}
+    callback_result = await custom_error_plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result=error_result,
+    )
+    self.assertIsNotNone(callback_result)
+    self.assertEqual(
+        callback_result["response_type"], REFLECT_AND_RETRY_RESPONSE_TYPE
+    )
+    self.assertEqual(callback_result["retry_count"], 1)
+
+    # Test success handling
+    success_result = {"status": "success", "data": "operation completed"}
+    callback_result = await custom_error_plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result=success_result,
+    )
+    self.assertIsNone(callback_result)
+
+  async def test_retry_state_management(self):
+    """Test retry state management with custom errors and mixed error types."""
+    custom_error_plugin = self.get_custom_error_plugin()
+    mock_tool = self.get_mock_tool()
+    mock_tool_context = self.get_mock_tool_context()
+    sample_tool_args = self.get_sample_tool_args()
+    custom_error_plugin.set_error_condition(
+        lambda result: result if result.get("failed") else None
+    )
+
+    # Custom error followed by exception
+    custom_error = {"failed": True, "reason": "Network timeout"}
+    result1 = await custom_error_plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result=custom_error,
+    )
+    self.assertEqual(result1["retry_count"], 1)
+
+    # Exception should increment retry count
+    exception = ValueError("Invalid parameter")
+    result2 = await custom_error_plugin.on_tool_error_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        error=exception,
+    )
+    self.assertEqual(result2["retry_count"], 2)
+
+    # Success should reset
+    success = {"result": "success"}
+    result3 = await custom_error_plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result=success,
+    )
+    self.assertIsNone(result3)
+
+    # Next error should start fresh
+    result4 = await custom_error_plugin.after_tool_callback(
+        tool=mock_tool,
+        tool_args=sample_tool_args,
+        tool_context=mock_tool_context,
+        result=custom_error,
+    )
+    self.assertEqual(result4["retry_count"], 1)
+
+  async def test_hallucinating_tool_name(self):
+    """Test that hallucinating tool name is handled correctly."""
+    wrong_function_call = types.Part.from_function_call(
+        name="increase_by_one", args={"x": 1}
+    )
+    correct_function_call = types.Part.from_function_call(
+        name="increase", args={"x": 1}
+    )
+    responses: list[types.Content] = [
+        wrong_function_call,
+        correct_function_call,
+        "response1",
+    ]
+    mock_model = testing_utils.MockModel.create(responses=responses)
+
+    function_called = 0
+
+    def increase(x: int) -> int:
+      nonlocal function_called
+      function_called += 1
+      return x + 1
+
+    agent = LlmAgent(name="root_agent", model=mock_model, tools=[increase])
+    runner = testing_utils.TestInMemoryRunner(
+        agent=agent, plugins=[self.get_plugin()]
+    )
+
+    events = await runner.run_async_with_new_session("test")
+    # Filter out agent_state events (no content).
+    events = [e for e in events if e.content is not None]
+
+    # Assert that the first event is a function call with the wrong name
+    assert events[0].content.parts[0].function_call.name == "increase_by_one"
+
+    # Assert that the second event is a function response with the
+    # reflection_guidance
+    assert (
+        events[1].content.parts[0].function_response.response["error_type"]
+        == "ValueError"
+    )
+    assert (
+        events[1].content.parts[0].function_response.response["retry_count"]
+        == 1
+    )
+    assert (
+        "Wrong Function Name"
+        in events[1]
+        .content.parts[0]
+        .function_response.response["reflection_guidance"]
+    )
+
+    # Assert that the third event is a function call with the correct name
+    assert events[2].content.parts[0].function_call.name == "increase"
+    self.assertEqual(function_called, 1)