feat: add new callbacks to handle tool and model errors

This CL add new callbacks in plugin system:
- `on_tool_error_callback`
- `on_model_error_callback`

This allow the user to create plugins that can handle errors.

PiperOrigin-RevId: 786469646

tests/unittests/flows/llm_flows/test_plugin_model_callbacks.py

@@ -20,19 +20,33 @@ from google.adk.models.llm_request import LlmRequest
 from google.adk.models.llm_response import LlmResponse
 from google.adk.plugins.base_plugin import BasePlugin
 from google.genai import types
+from google.genai.errors import ClientError
 import pytest
 
 from ... import testing_utils
 
+mock_error = ClientError(
+    code=429,
+    response_json={
+        'error': {
+            'code': 429,
+            'message': 'Quota exceeded.',
+            'status': 'RESOURCE_EXHAUSTED',
+        }
+    },
+)
+
 
 class MockPlugin(BasePlugin):
   before_model_text = 'before_model_text from MockPlugin'
   after_model_text = 'after_model_text from MockPlugin'
+  on_model_error_text = 'on_model_error_text from MockPlugin'
 
   def __init__(self, name='mock_plugin'):
     self.name = name
     self.enable_before_model_callback = False
     self.enable_after_model_callback = False
+    self.enable_on_model_error_callback = False
     self.before_model_response = LlmResponse(
         content=testing_utils.ModelContent(
             [types.Part.from_text(text=self.before_model_text)]
@@ -43,6 +57,11 @@ class MockPlugin(BasePlugin):
             [types.Part.from_text(text=self.after_model_text)]
         )
     )
+    self.on_model_error_response = LlmResponse(
+        content=testing_utils.ModelContent(
+            [types.Part.from_text(text=self.on_model_error_text)]
+        )
+    )
 
   async def before_model_callback(
       self, *, callback_context: CallbackContext, llm_request: LlmRequest
@@ -58,6 +77,17 @@ class MockPlugin(BasePlugin):
       return None
     return self.after_model_response
 
+  async def on_model_error_callback(
+      self,
+      *,
+      callback_context: CallbackContext,
+      llm_request: LlmRequest,
+      error: Exception,
+  ) -> Optional[LlmResponse]:
+    if not self.enable_on_model_error_callback:
+      return None
+    return self.on_model_error_response
+
 
 CANONICAL_MODEL_CALLBACK_CONTENT = 'canonical_model_callback_content'
 
@@ -124,5 +154,36 @@ def test_before_model_callback_fallback_model(mock_plugin):
   ]
 
 
+def test_on_model_error_callback_with_plugin(mock_plugin):
+  """Tests that the model error is handled by the plugin."""
+  mock_model = testing_utils.MockModel.create(error=mock_error, responses=[])
+  mock_plugin.enable_on_model_error_callback = True
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+  )
+
+  runner = testing_utils.InMemoryRunner(agent, plugins=[mock_plugin])
+
+  assert testing_utils.simplify_events(runner.run('test')) == [
+      ('root_agent', mock_plugin.on_model_error_text),
+  ]
+
+
+def test_on_model_error_callback_fallback_to_runner(mock_plugin):
+  """Tests that the model error is not handled and falls back to raise from runner."""
+  mock_model = testing_utils.MockModel.create(error=mock_error, responses=[])
+  mock_plugin.enable_on_model_error_callback = False
+  agent = Agent(
+      name='root_agent',
+      model=mock_model,
+  )
+
+  try:
+    testing_utils.InMemoryRunner(agent, plugins=[mock_plugin])
+  except Exception as e:
+    assert e == mock_error
+
+
 if __name__ == '__main__':
   pytest.main([__file__])

tests/unittests/flows/llm_flows/test_plugin_tool_callbacks.py

@@ -24,19 +24,35 @@ from google.adk.tools.base_tool import BaseTool
 from google.adk.tools.function_tool import FunctionTool
 from google.adk.tools.tool_context import ToolContext
 from google.genai import types
+from google.genai.errors import ClientError
 import pytest
 
 from ... import testing_utils
 
+mock_error = ClientError(
+    code=429,
+    response_json={
+        "error": {
+            "code": 429,
+            "message": "Quota exceeded.",
+            "status": "RESOURCE_EXHAUSTED",
+        }
+    },
+)
+
 
 class MockPlugin(BasePlugin):
   before_tool_response = {"MockPlugin": "before_tool_response from MockPlugin"}
   after_tool_response = {"MockPlugin": "after_tool_response from MockPlugin"}
+  on_tool_error_response = {
+      "MockPlugin": "on_tool_error_response from MockPlugin"
+  }
 
   def __init__(self, name="mock_plugin"):
     self.name = name
     self.enable_before_tool_callback = False
     self.enable_after_tool_callback = False
+    self.enable_on_tool_error_callback = False
 
   async def before_tool_callback(
       self,
@@ -61,6 +77,18 @@ class MockPlugin(BasePlugin):
       return None
     return self.after_tool_response
 
+  async def on_tool_error_callback(
+      self,
+      *,
+      tool: BaseTool,
+      tool_args: dict[str, Any],
+      tool_context: ToolContext,
+      error: Exception,
+  ) -> Optional[dict]:
+    if not self.enable_on_tool_error_callback:
+      return None
+    return self.on_tool_error_response
+
 
 @pytest.fixture
 def mock_tool():
@@ -70,6 +98,14 @@ def mock_tool():
   return FunctionTool(simple_fn)
 
 
+@pytest.fixture
+def mock_error_tool():
+  def raise_error_fn(**kwargs) -> Dict[str, Any]:
+    raise mock_error
+
+  return FunctionTool(raise_error_fn)
+
+
 @pytest.fixture
 def mock_plugin():
   return MockPlugin()
@@ -124,5 +160,30 @@ async def test_async_after_tool_callback(mock_tool, mock_plugin):
   assert part.function_response.response == mock_plugin.after_tool_response
 
 
+@pytest.mark.asyncio
+async def test_async_on_tool_error_use_plugin_response(
+    mock_error_tool, mock_plugin
+):
+  mock_plugin.enable_on_tool_error_callback = True
+
+  result_event = await invoke_tool_with_plugin(mock_error_tool, mock_plugin)
+
+  assert result_event is not None
+  part = result_event.content.parts[0]
+  assert part.function_response.response == mock_plugin.on_tool_error_response
+
+
+@pytest.mark.asyncio
+async def test_async_on_tool_error_fallback_to_runner(
+    mock_error_tool, mock_plugin
+):
+  mock_plugin.enable_on_tool_error_callback = False
+
+  try:
+    await invoke_tool_with_plugin(mock_error_tool, mock_plugin)
+  except Exception as e:
+    assert e == mock_error
+
+
 if __name__ == "__main__":
   pytest.main([__file__])