test_async_inner_command_retry_safe() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_wrap_model_call_state_update.py lines 735–779

    async def test_async_inner_command_retry_safe(self) -> None:
        """Async: when outer retries, only last inner command is used."""
        call_count = 0

        class MyState(AgentState):
            attempt: str

        class OuterMiddleware(AgentMiddleware):
            async def awrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
            ) -> ModelResponse:
                # Call handler twice (simulating retry)
                await handler(request)
                return await handler(request)

        class InnerMiddleware(AgentMiddleware):
            state_schema = MyState  # type: ignore[assignment]

            async def awrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
            ) -> ExtendedModelResponse:
                nonlocal call_count
                call_count += 1
                response = await handler(request)
                return ExtendedModelResponse(
                    model_response=response,
                    command=Command(update={"attempt": f"attempt_{call_count}"}),
                )

        model = GenericFakeChatModel(
            messages=iter([AIMessage(content="First"), AIMessage(content="Second")])
        )
        agent = create_agent(
            model=model,
            middleware=[OuterMiddleware(), InnerMiddleware()],
        )

        result = await agent.ainvoke({"messages": [HumanMessage("Hi")]})

        messages = result["messages"]
        assert any(m.content == "Second" for m in messages)