From 4eaab0284a621b5d19b557807e018ae640f4c352 Mon Sep 17 00:00:00 2001
From: Caspar Broekhuizen <casparbroekhuizen@gmail.com>
Date: Tue, 26 Aug 2025 15:49:55 -0700
Subject: [PATCH] feat(prebuilt): add support for handle_tool_errors with
 single exception type (#6021)

This PR extends the `ToolNode` `handle_tool_errors` property to also
accept a single exception type: `type[Exception]`.

Example:
```python
tool_node_with_error_handling = ToolNode(
    [tool1, tool2, tool3], handle_tool_errors=ValueError
).invoke(...)
"""A tool node that will only retry a failed tool call if a ValueError is raised"""
```
---
 libs/prebuilt/langgraph/prebuilt/tool_node.py | 36 ++++++++++++++-----
 libs/prebuilt/tests/test_tool_node.py         |  2 ++
 2 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/libs/prebuilt/langgraph/prebuilt/tool_node.py b/libs/prebuilt/langgraph/prebuilt/tool_node.py
index 122c71b97..bea4434b3 100644
--- a/libs/prebuilt/langgraph/prebuilt/tool_node.py
+++ b/libs/prebuilt/langgraph/prebuilt/tool_node.py
@@ -157,6 +157,7 @@ def _handle_tool_error(
         bool,
         str,
         Callable[..., str],
+        type[Exception],
         tuple[type[Exception], ...],
     ],
 ) -> str:
@@ -184,12 +185,14 @@ def _handle_tool_error(
         The tuple case is handled by the caller through exception type checking,
         not by this function directly.
     """
-    if isinstance(flag, (bool, tuple)):
+    if isinstance(flag, (bool, tuple)) or (
+        isinstance(flag, type) and issubclass(flag, Exception)
+    ):
         content = TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e))
     elif isinstance(flag, str):
         content = flag
     elif callable(flag):
-        content = flag(e)
+        content = flag(e)  # type: ignore [assignment, call-arg]
     else:
         raise ValueError(
             f"Got unexpected type of `handle_tool_error`. Expected bool, str "
@@ -309,6 +312,7 @@ class ToolNode(RunnableCallable):
               error template containing the exception details.
             - **str**: Catch all errors and return a ToolMessage with this custom
               error message string.
+            - **type[Exception]**: Only catch exceptions with the specified type and return the default error message for it.
             - **tuple[type[Exception], ...]**: Only catch exceptions with the specified
               types and return default error messages for them.
             - **Callable[..., str]**: Catch exceptions matching the callable's signature
@@ -373,7 +377,7 @@ class ToolNode(RunnableCallable):
         name: str = "tools",
         tags: Optional[list[str]] = None,
         handle_tool_errors: Union[
-            bool, str, Callable[..., str], tuple[type[Exception], ...]
+            bool, str, Callable[..., str], type[Exception], tuple[type[Exception], ...]
         ] = _default_handle_tool_errors,
         messages_key: str = "messages",
     ) -> None:
@@ -518,9 +522,16 @@ class ToolNode(RunnableCallable):
         except GraphBubbleUp as e:
             raise e
         except Exception as e:
-            if isinstance(self._handle_tool_errors, tuple):
-                handled_types: tuple = self._handle_tool_errors
-            elif callable(self._handle_tool_errors):
+            handled_types: tuple[type[Exception], ...]
+            if isinstance(self._handle_tool_errors, type) and issubclass(
+                self._handle_tool_errors, Exception
+            ):
+                handled_types = (self._handle_tool_errors,)
+            elif isinstance(self._handle_tool_errors, tuple):
+                handled_types = self._handle_tool_errors
+            elif callable(self._handle_tool_errors) and not isinstance(
+                self._handle_tool_errors, type
+            ):
                 handled_types = _infer_handled_types(self._handle_tool_errors)
             else:
                 # default behavior is catching all exceptions
@@ -580,9 +591,16 @@ class ToolNode(RunnableCallable):
         except GraphBubbleUp as e:
             raise e
         except Exception as e:
-            if isinstance(self._handle_tool_errors, tuple):
-                handled_types: tuple = self._handle_tool_errors
-            elif callable(self._handle_tool_errors):
+            handled_types: tuple[type[Exception], ...]
+            if isinstance(self._handle_tool_errors, type) and issubclass(
+                self._handle_tool_errors, Exception
+            ):
+                handled_types = (self._handle_tool_errors,)
+            elif isinstance(self._handle_tool_errors, tuple):
+                handled_types = self._handle_tool_errors
+            elif callable(self._handle_tool_errors) and not isinstance(
+                self._handle_tool_errors, type
+            ):
                 handled_types = _infer_handled_types(self._handle_tool_errors)
             else:
                 # default behavior is catching all exceptions
diff --git a/libs/prebuilt/tests/test_tool_node.py b/libs/prebuilt/tests/test_tool_node.py
index ce59baaaf..a26b9d136 100644
--- a/libs/prebuilt/tests/test_tool_node.py
+++ b/libs/prebuilt/tests/test_tool_node.py
@@ -274,10 +274,12 @@ async def test_tool_node_error_handling() -> None:
 
     # test catching all exceptions, via:
     # - handle_tool_errors = True
+    # - passing a single exception
     # - passing a tuple of all exceptions
     # - passing a callable with all exceptions in the signature
     for handle_tool_errors in (
         True,
+        Exception,
         (ValueError, ToolException, ToolInvocationError),
         handle_all,
     ):