Let HookToolset retry a failed tool call instead of failing the task

providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py

@@ -24,6 +24,7 @@
 import types
 from typing import TYPE_CHECKING, Any, Union, get_args, get_origin, get_type_hints
 
+from pydantic_ai.exceptions import ModelRetry
 from pydantic_ai.tools import ToolDefinition
 from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool
 
@@ -61,6 +62,10 @@ class HookToolset(AbstractToolset[Any]):
         auto-discovery is intentionally not supported for safety.
     :param tool_name_prefix: Optional prefix prepended to each tool name
         (e.g. ``"s3_"`` → ``"s3_list_keys"``).
+
+    When a tool call fails, the hook's own error message is returned to the agent as a
+    retry (:class:`pydantic_ai.ModelRetry`) so the model can correct its arguments within
+    the run instead of failing the task.
     """
 
     def __init__(
@@ -136,7 +141,10 @@ async def call_tool(
     ) -> Any:
         method_name = name.removeprefix(self._tool_name_prefix) if self._tool_name_prefix else name
         method: Callable[..., Any] = getattr(self._hook, method_name)
-        result = method(**tool_args)
+        try:
+            result = method(**tool_args)
+        except Exception as e:
+            raise ModelRetry(f"The {name} tool failed: {e}") from e
         return _serialize_for_llm(result)
 
 

providers/common/ai/tests/unit/common/ai/toolsets/test_hook.py

@@ -20,6 +20,7 @@
 from unittest.mock import MagicMock
 
 import pytest
+from pydantic_ai.exceptions import ModelRetry
 from pydantic_core import ValidationError
 
 from airflow.providers.common.ai.toolsets.hook import (
@@ -50,6 +51,10 @@ def read_file(self, key: str) -> str:
     def no_docstring(self, x: int) -> int:
         return x * 2
 
+    def failing_method(self, x: int) -> int:
+        """Always raise, for testing call_tool failure handling."""
+        raise RuntimeError("boom")
+
     def request(
         self, endpoint: str | None = None, data: dict[str, object] | str | None = None, **kwargs: object
     ) -> dict[str, object]:
@@ -203,6 +208,21 @@ def test_dispatches_with_prefix(self):
         )
         assert result == "contents of test.txt"
 
+    def test_wraps_failure_in_model_retry(self):
+        hook = _FakeHook()
+        ts = HookToolset(hook, allowed_methods=["failing_method"])
+        tools = asyncio.run(ts.get_tools(ctx=MagicMock()))
+
+        with pytest.raises(ModelRetry, match="failing_method tool failed: boom"):
+            asyncio.run(
+                ts.call_tool(
+                    "failing_method",
+                    {"x": 1},
+                    ctx=MagicMock(),
+                    tool=tools["failing_method"],
+                )
+            )
+
 
 class TestBuildJsonSchemaFromSignature:
     def test_basic_types(self):