From 27160cd1fe8ab2efef448a423b4fa543d9cf1ca1 Mon Sep 17 00:00:00 2001 From: Chester Curme Date: Fri, 10 Jul 2026 11:14:33 -0400 Subject: [PATCH] support async on_error --- .../langchain/agents/middleware/tool_error.py | 50 +++++++++++++------ .../implementations/test_tool_error.py | 48 ++++++++++++++++++ 2 files changed, 84 insertions(+), 14 deletions(-) diff --git a/libs/langchain_v1/langchain/agents/middleware/tool_error.py b/libs/langchain_v1/langchain/agents/middleware/tool_error.py index ca2df88624d..3a307b8f94e 100644 --- a/libs/langchain_v1/langchain/agents/middleware/tool_error.py +++ b/libs/langchain_v1/langchain/agents/middleware/tool_error.py @@ -2,6 +2,7 @@ from __future__ import annotations +import inspect from collections.abc import Callable from typing import TYPE_CHECKING, Any @@ -21,8 +22,14 @@ if TYPE_CHECKING: Catch = tuple[type[Exception], ...] | Callable[[Exception], bool] """Exceptions to catch: a tuple of exception types, or a predicate ``(exc) -> bool``.""" -OnError = Callable[[Exception, "ToolCallRequest"], "str | list[str | dict[Any, Any]]"] -"""Formatter for a caught exception, returning `ToolMessage` content.""" +OnError = Callable[ + [Exception, "ToolCallRequest"], + "str | list[str | dict[Any, Any]] | Awaitable[str | list[str | dict[Any, Any]]]", +] +"""Handler for a caught exception, returning `ToolMessage` content. + +May be sync or async. An async `on_error` requires async execution (`ainvoke`/`astream`). +""" def _should_catch(exc: Exception, catch: Catch) -> bool: @@ -130,6 +137,31 @@ class ToolErrorMiddleware(AgentMiddleware[AgentState[ResponseT], ContextT, Respo """ return f"Tool '{tool_name}' failed with {type(exc).__name__}." + def _sync_content( + self, exc: Exception, request: ToolCallRequest, tool_name: str + ) -> str | list[str | dict[Any, Any]]: + """Resolve the error message content in a sync context.""" + if self.on_error is None: + return self._format_error(tool_name, exc) + result = self.on_error(exc, request) + if inspect.isawaitable(result): + if inspect.iscoroutine(result): + result.close() + msg = "async on_error requires async execution (ainvoke or astream)" + raise TypeError(msg) + return result + + async def _async_content( + self, exc: Exception, request: ToolCallRequest, tool_name: str + ) -> str | list[str | dict[Any, Any]]: + """Resolve the error message content in an async context.""" + if self.on_error is None: + return self._format_error(tool_name, exc) + result = self.on_error(exc, request) + if inspect.isawaitable(result): + return await result + return result + def wrap_tool_call( self, request: ToolCallRequest, @@ -157,13 +189,8 @@ class ToolErrorMiddleware(AgentMiddleware[AgentState[ResponseT], ContextT, Respo except Exception as exc: if not _should_catch(exc, self.catch): raise - content = ( - self.on_error(exc, request) - if self.on_error is not None - else self._format_error(tool_name, exc) - ) return ToolMessage( - content=content, + content=self._sync_content(exc, request, tool_name), tool_call_id=request.tool_call["id"], name=tool_name, status="error", @@ -188,13 +215,8 @@ class ToolErrorMiddleware(AgentMiddleware[AgentState[ResponseT], ContextT, Respo except Exception as exc: if not _should_catch(exc, self.catch): raise - content = ( - self.on_error(exc, request) - if self.on_error is not None - else self._format_error(tool_name, exc) - ) return ToolMessage( - content=content, + content=await self._async_content(exc, request, tool_name), tool_call_id=request.tool_call["id"], name=tool_name, status="error", diff --git a/libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_tool_error.py b/libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_tool_error.py index 89e2f26ad17..dc03c5a575d 100644 --- a/libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_tool_error.py +++ b/libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_tool_error.py @@ -3,6 +3,7 @@ import pytest from langchain_core.messages import HumanMessage, ToolCall, ToolMessage from langchain_core.tools import tool +from langgraph.prebuilt.tool_node import ToolCallRequest from langchain.agents.factory import create_agent from langchain.agents.middleware import ToolErrorMiddleware @@ -54,3 +55,50 @@ def test_tool_error_uncaught_propagates() -> None: with pytest.raises(ValueError, match="secret detail"): agent.invoke({"messages": [HumanMessage("go")]}) + + +def test_tool_error_on_error_formats_content() -> None: + """`on_error` sets the ToolMessage content and can use the tool call context.""" + + def on_error(exc: Exception, request: ToolCallRequest) -> str: + return ( + f"`{request.tool_call['name']}` raised {type(exc).__name__} for " + f"{request.tool_call['args']}; fix the arguments and retry." + ) + + agent = create_agent( + model=_model(), + tools=[failing_tool], + middleware=[ToolErrorMiddleware(catch=(ValueError,), on_error=on_error)], + ) + + result = agent.invoke({"messages": [HumanMessage("go")]}) + + tool_messages = [m for m in result["messages"] if isinstance(m, ToolMessage)] + assert len(tool_messages) == 1 + assert tool_messages[0].status == "error" + assert tool_messages[0].content == ( + "`failing_tool` raised ValueError for {'value': 'x'}; fix the arguments and retry." + ) + # Custom formatter controls disclosure — the raw exception message is not leaked. + assert "secret detail" not in tool_messages[0].content + + +async def test_tool_error_async_on_error() -> None: + """An async `on_error` is awaited under async execution.""" + + async def on_error(exc: Exception, request: ToolCallRequest) -> str: + return f"async handled `{request.tool_call['name']}`: {type(exc).__name__}" + + agent = create_agent( + model=_model(), + tools=[failing_tool], + middleware=[ToolErrorMiddleware(catch=(ValueError,), on_error=on_error)], + ) + + result = await agent.ainvoke({"messages": [HumanMessage("go")]}) + + tool_messages = [m for m in result["messages"] if isinstance(m, ToolMessage)] + assert len(tool_messages) == 1 + assert tool_messages[0].status == "error" + assert tool_messages[0].content == "async handled `failing_tool`: ValueError"