test_tools_preserved_across_multiple_middleware() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_tools.py lines 247–318

def test_tools_preserved_across_multiple_middleware() -> None:
    """Test that tool modifications by one middleware are visible to the next."""
    modification_order: list[list[str]] = []

    @tool
    def tool_a(value: str) -> str:
        """Tool A."""
        return "A"

    @tool
    def tool_b(value: str) -> str:
        """Tool B."""
        return "B"

    @tool
    def tool_c(value: str) -> str:
        """Tool C."""
        return "C"

    class FirstMiddleware(AgentMiddleware):
        def wrap_model_call(
            self,
            request: ModelRequest,
            handler: Callable[[ModelRequest], ModelResponse],
        ) -> ModelCallResult:
            tools: list[str] = []
            filtered_tools: list[BaseTool | dict[str, Any]] = []
            for t in request.tools:
                assert isinstance(t, BaseTool)
                tools.append(t.name)
                # Remove tool_c
                if t.name != "tool_c":
                    filtered_tools.append(t)
            modification_order.append(tools)
            request = request.override(tools=filtered_tools)
            return handler(request)

    class SecondMiddleware(AgentMiddleware):
        def wrap_model_call(
            self,
            request: ModelRequest,
            handler: Callable[[ModelRequest], ModelResponse],
        ) -> ModelCallResult:
            tools: list[str] = []
            filtered_tools: list[BaseTool | dict[str, Any]] = []
            for t in request.tools:
                assert isinstance(t, BaseTool)
                # Should not see tool_c here
                assert t.name != "tool_c"
                tools.append(t.name)
                # Remove tool_b
                if t.name != "tool_b":
                    filtered_tools.append(t)
            modification_order.append(tools)
            request = request.override(tools=filtered_tools)
            return handler(request)

    agent = create_agent(
        model=FakeToolCallingModel(),
        tools=[tool_a, tool_b, tool_c],
        system_prompt="You are a helpful assistant.",
        middleware=[FirstMiddleware(), SecondMiddleware()],
    )

    agent.invoke({"messages": [HumanMessage("Hello")]})

    # Verify the modification sequence
    assert len(modification_order) == 2
    # First middleware sees all three tools
    assert set(modification_order[0]) == {"tool_a", "tool_b", "tool_c"}
    # Second middleware sees tool_c removed
    assert set(modification_order[1]) == {"tool_a", "tool_b"}