So I create the agent with tools using langchain create_agent method, after the execution of the certain tool I would like to end the agent’s flow unless the exception in this tool have been made.
I use wrap_tool_call example from docs:
@wrap_tool_call
async def handle_tool_errors(request, handler):
"""Handle tool execution errors with custom messages."""
try:
return await handler(request)
except Exception as thrown:
# Return a custom error message to the model
handler.return_direct = False
return ToolMessage(
content=f"Tool error: Please check your input and try again. ({str(thrown)})",
tool_call_id=request.tool_call["id"]
The problem is that the flow is still finishing even with the exception being handled in wrap_tool_call
So, what is the best practice for this? How to force the agent to finish after only successful execution of the tool?
With create_agent, the graph exits after tool execution if the tool(s) called have return_direct=True - regardless of whether the tool returned an error. Flipping handler.return_direct = False in a wrapper doesn’t affect routing.
Hmmm maybe don’t mark the tool return_direct=True. Instead, keep return_direct=False (default) and use your wrap_tool_call to:
return an error ToolMessage on exceptions so the agent continues the loop, and
return a Command(goto=END) only on success to finish immediately.
Try sth like this (or let me know if you need sth else)
Keep the tool as @tool() (no return_direct).
Use a middleware wrapper that:
returns an error ToolMessage on exception (so the LLM can repair), and
returns a Command(goto=END) on success to end immediately.
from langchain.agents.middleware import wrap_tool_call
from langchain_core.messages import ToolMessage
from langgraph.types import Command
from langgraph.constants import END
@wrap_tool_call
async def finish_on_success(req, handler):
try:
result = await handler(req) # ToolMessage or Command
except Exception as e:
return ToolMessage(
content=f"Tool error: Please check your input and try again. ({e})",
tool_call_id=req.tool_call["id"],
name=req.tool_call["name"],
status="error",
)
# If tool itself returned a Command, pass it through
if isinstance(result, dict) and ("goto" in result or "update" in result):
return result
# End only on success; continue otherwise
if isinstance(result, ToolMessage) and getattr(result, "status", None) != "error":
return Command(update={"messages": [result]}, goto=END)
return result # error ToolMessage → continue loop
Yeah I tried this approach with finishing execution via wrapper, but for some reason it doesn’t seem to be finishing when Command(goto=END) is returned. I tried using it like this (tool itself returns hard-coded string), I did also double check if condition actually applies and it does seem to be updating messages in the state, but the goto=END doesnt seem to be doing anything:
@wrap_tool_call
async def handle_tool_errors(request, handler):
"""Handle tool execution errors with custom messages."""
try:
result = await handler(request)
if result.content == "Response submitted":
return Command(update={"messages": [result]}, goto=END)
else:
return result
except Exception as thrown:
# Return a custom error message to the model
handler.return_direct = False
return ToolMessage(
content=f"Tool error: Please check your input and try again. ({str(thrown)})",
tool_call_id=request.tool_call["id"]
)
Hmm, I decided to lookup create_agent method and how it interacts with “return_direct=True” tools, and it seems that it only includes exit_node in tools destinations if there is any tool with return_direct parameter?
Pycharm’s AI Agent Debugger also seems to confirm this on graph
