fix(prebuilt): handle asyncio.CancelledError in ToolNode

[Saakshi Gupta (saakshigupta2002)]

Pull Request: fix(prebuilt): handle asyncio.CancelledError in ToolNode

Summary

This PR fixes issue #6726 where asyncio.CancelledError is not caught by the ToolNode error handling mechanism, even when handle_tool_errors=True is set.

Problem

When a tool execution is cancelled via asyncio.CancelledError, the ToolNode does not create an error ToolMessage. This leaves the message history in an invalid state where an AIMessage has tool_calls without corresponding ToolMessages, causing INVALID_CHAT_HISTORY errors on subsequent LLM calls.

Root Cause

asyncio.CancelledError inherits from BaseException, not Exception. The existing error handling in ToolNode._execute_tool_async() and _arun_one() uses except Exception as e:, which doesn't catch CancelledError.

Solution

• Added explicit except asyncio.CancelledError: handlers before the except Exception: blocks
• When handle_tool_errors=True, a ToolMessage with status="error" is returned
• When handle_tool_errors=False, the exception is re-raised as expected
• Added a new constant TOOL_CANCELLED_ERROR_TEMPLATE for consistent error messages

Changes

libs/prebuilt/langgraph/prebuilt/tool_node.py

1. Added constant (line 118):

   TOOL_CANCELLED_ERROR_TEMPLATE = "Tool execution was cancelled."

2. Modified _execute_tool_async - Added CancelledError handler before GraphBubbleUp:

   except asyncio.CancelledError:
       if self._handle_tool_errors:
           return ToolMessage(
               content=TOOL_CANCELLED_ERROR_TEMPLATE,
               name=call["name"],
               tool_call_id=call["id"],
               status="error",
           )
       raise

3. Modified _arun_one - Added CancelledError handler for wrapper cancellation:

   except asyncio.CancelledError:
       if self._handle_tool_errors:
           return ToolMessage(
               content=TOOL_CANCELLED_ERROR_TEMPLATE,
               name=tool_request.tool_call["name"],
               tool_call_id=tool_request.tool_call["id"],
               status="error",
           )
       raise

libs/prebuilt/tests/test_tool_node.py

Added three test cases:

1. test_tool_node_cancelled_error_handled - Verifies that a ToolMessage with error status is returned when handle_tool_errors=True and a tool is cancelled

2. test_tool_node_cancelled_error_not_handled - Verifies that CancelledError is raised when handle_tool_errors=False

3. test_tool_node_cancelled_error_in_wrapper - Verifies that CancelledError raised in awrap_tool_call is handled properly

Test Plan

• Run existing ToolNode tests: pytest libs/prebuilt/tests/test_tool_node.py -v
• Run new cancellation tests: pytest libs/prebuilt/tests/test_tool_node.py -v -k cancelled
• Verify the reproduction script from issue handle_tool_errors=True does not catch asyncio.CancelledError, causing INVALID_CHAT_HISTORY #6726 now works correctly
• Test with real async tool that gets cancelled (e.g., timeout scenario)

Reproduction Script (from issue)

import asyncio
from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.tools import tool
from langgraph.prebuilt import ToolNode

@tool
async def slow_tool(query: str) -> str:
    """A tool that takes time to complete."""
    await asyncio.sleep(10)
    return f"Result: {query}"

async def test_cancelled_error():
    tool_node = ToolNode(tools=[slow_tool], handle_tool_errors=True)

    ai_message = AIMessage(
        content="",
        tool_calls=[{"id": "call_1", "name": "slow_tool", "args": {"query": "test"}}]
    )
    state = {"messages": [HumanMessage(content="test"), ai_message]}
    config = {"configurable": {"thread_id": "test"}}

    task = asyncio.create_task(tool_node.ainvoke(state, config))
    await asyncio.sleep(0.1)
    task.cancel()

    try:
        result = await task
        print("Got result:", result)  # Now returns ToolMessage with error status
    except asyncio.CancelledError:
        print("CancelledError raised - no ToolMessage created!")  # This no longer happens

asyncio.run(test_cancelled_error())

Before/After Comparison

Scenario	Before	After
handle_tool_errors=True, tool cancelled	CancelledError raised, no ToolMessage	ToolMessage with status="error" returned
handle_tool_errors=False, tool cancelled	CancelledError raised	CancelledError raised (unchanged)
Message history validity	Invalid (AIMessage with tool_calls, no ToolMessage)	Valid (matching ToolMessage created)

Related Issues

• Fixes handle_tool_errors=True does not catch asyncio.CancelledError, causing INVALID_CHAT_HISTORY #6726

Checklist

• Code follows the project's coding style
• Changes are covered by tests
• All existing tests pass
• Documentation is updated (docstrings in code)
• Commit message follows conventional commits format