test_model_retry_multiple_middleware_composition() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_model_retry.py lines 662–698

def test_model_retry_multiple_middleware_composition() -> None:
    """Test ModelRetryMiddleware composes correctly with other middleware."""
    call_log = []

    # Custom middleware that logs calls
    @wrap_model_call
    def logging_middleware(
        request: ModelRequest,
        handler: Callable[[ModelRequest], ModelResponse],
    ) -> ModelCallResult:
        call_log.append("before_model")
        response = handler(request)
        call_log.append("after_model")
        return response

    model = FakeToolCallingModel()

    retry = ModelRetryMiddleware(max_retries=2, initial_delay=0.01, jitter=False)

    agent = create_agent(
        model=model,
        tools=[],
        middleware=[logging_middleware, retry],
        checkpointer=InMemorySaver(),
    )

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

    # Both middleware should be called
    assert call_log == ["before_model", "after_model"]

    ai_messages = [m for m in result["messages"] if isinstance(m, AIMessage)]
    assert len(ai_messages) >= 1
    assert "Hello" in ai_messages[-1].content