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/plugins/test_base_plugin.py

@@ -67,12 +67,18 @@ class FullOverridePlugin(BasePlugin):
   async def after_tool_callback(self, **kwargs) -> str:
     return "overridden_after_tool"
 
+  async def on_tool_error_callback(self, **kwargs) -> str:
+    return "overridden_on_tool_error"
+
   async def before_model_callback(self, **kwargs) -> str:
     return "overridden_before_model"
 
   async def after_model_callback(self, **kwargs) -> str:
     return "overridden_after_model"
 
+  async def on_model_error_callback(self, **kwargs) -> str:
+    return "overridden_on_model_error"
+
 
 def test_base_plugin_initialization():
   """Tests that a plugin is initialized with the correct name."""
@@ -137,6 +143,15 @@ async def test_base_plugin_default_callbacks_return_none():
       )
       is None
   )
+  assert (
+      await plugin.on_tool_error_callback(
+          tool=mock_context,
+          tool_args={},
+          tool_context=mock_context,
+          error=Exception(),
+      )
+      is None
+  )
   assert (
       await plugin.before_model_callback(
           callback_context=mock_context, llm_request=mock_context
@@ -149,6 +164,14 @@ async def test_base_plugin_default_callbacks_return_none():
       )
       is None
   )
+  assert (
+      await plugin.on_model_error_callback(
+          callback_context=mock_context,
+          llm_request=mock_context,
+          error=Exception(),
+      )
+      is None
+  )
 
 
 @pytest.mark.asyncio
@@ -170,6 +193,7 @@ async def test_base_plugin_all_callbacks_can_be_overridden():
   mock_llm_request = Mock(spec=LlmRequest)
   mock_llm_response = Mock(spec=LlmResponse)
   mock_event = Mock(spec=Event)
+  mock_error = Mock(spec=Exception)
 
   # Call each method and assert it returns the unique string from the override.
   # This proves that the subclass's method was executed.
@@ -237,3 +261,20 @@ async def test_base_plugin_all_callbacks_can_be_overridden():
       )
       == "overridden_after_tool"
   )
+  assert (
+      await plugin.on_tool_error_callback(
+          tool=mock_tool,
+          tool_args={},
+          tool_context=mock_tool_context,
+          error=mock_error,
+      )
+      == "overridden_on_tool_error"
+  )
+  assert (
+      await plugin.on_model_error_callback(
+          callback_context=mock_callback_context,
+          llm_request=mock_llm_request,
+          error=mock_error,
+      )
+      == "overridden_on_model_error"
+  )

tests/unittests/plugins/test_plugin_manager.py

@@ -77,12 +77,18 @@ class TestPlugin(BasePlugin):
   async def after_tool_callback(self, **kwargs):
     return await self._handle_callback("after_tool_callback")
 
+  async def on_tool_error_callback(self, **kwargs):
+    return await self._handle_callback("on_tool_error_callback")
+
   async def before_model_callback(self, **kwargs):
     return await self._handle_callback("before_model_callback")
 
   async def after_model_callback(self, **kwargs):
     return await self._handle_callback("after_model_callback")
 
+  async def on_model_error_callback(self, **kwargs):
+    return await self._handle_callback("on_model_error_callback")
+
 
 @pytest.fixture
 def service() -> PluginManager:
@@ -227,12 +233,23 @@ async def test_all_callbacks_are_supported(
   await service.run_after_tool_callback(
       tool=mock_context, tool_args={}, tool_context=mock_context, result={}
   )
+  await service.run_on_tool_error_callback(
+      tool=mock_context,
+      tool_args={},
+      tool_context=mock_context,
+      error=mock_context,
+  )
   await service.run_before_model_callback(
       callback_context=mock_context, llm_request=mock_context
   )
   await service.run_after_model_callback(
       callback_context=mock_context, llm_response=mock_context
   )
+  await service.run_on_model_error_callback(
+      callback_context=mock_context,
+      llm_request=mock_context,
+      error=mock_context,
+  )
 
   # Verify all callbacks were logged
   expected_callbacks = [
@@ -244,7 +261,9 @@ async def test_all_callbacks_are_supported(
       "after_agent_callback",
       "before_tool_callback",
       "after_tool_callback",
+      "on_tool_error_callback",
       "before_model_callback",
       "after_model_callback",
+      "on_model_error_callback",
   ]
   assert set(plugin1.call_log) == set(expected_callbacks)