test_wrap_tool_call_retry_on_error() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_wrap_tool_call.py lines 228–285

def test_wrap_tool_call_retry_on_error() -> None:
    """Test retry logic with wrap_tool_call decorator on failing tool."""
    attempt_counts = []

    @wrap_tool_call
    def retry_middleware(
        request: ToolCallRequest, handler: Callable[[ToolCallRequest], ToolMessage | Command[Any]]
    ) -> ToolMessage | Command[Any]:
        max_retries = 3
        last_error = None
        for attempt in range(max_retries):
            attempt_counts.append(attempt)
            try:
                return handler(request)
            except Exception as e:
                last_error = e
                if attempt == max_retries - 1:
                    # Return error message instead of raising
                    return ToolMessage(
                        content=f"Error after {max_retries} attempts: {last_error}",
                        tool_call_id=request.tool_call["id"],
                        name=request.tool_call["name"],
                        status="error",
                    )
                # Continue to retry
        # This line should never be reached due to return above
        return ToolMessage(
            content=f"Unexpected error: {last_error}",
            tool_call_id=request.tool_call["id"],
            name=request.tool_call["name"],
            status="error",
        )

    model = FakeToolCallingModel(
        tool_calls=[
            [ToolCall(name="failing_tool", args={"value": "test"}, id="1")],
            [],
        ]
    )

    agent = create_agent(
        model=model,
        tools=[failing_tool],
        middleware=[retry_middleware],
        checkpointer=InMemorySaver(),
    )

    result = agent.invoke(
        {"messages": [HumanMessage("Use failing tool")]},
        {"configurable": {"thread_id": "test"}},
    )

    # Should attempt 3 times before giving up
    assert len(attempt_counts) == 3
    assert attempt_counts == [0, 1, 2]
    tool_messages = [m for m in result["messages"] if isinstance(m, ToolMessage)]
    assert len(tool_messages) == 1
    assert "Error after 3 attempts" in tool_messages[0].content