test_inner_state_preserved_when_outer_has_no_conflict() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_wrap_model_call_state_update.py lines 486–531

    def test_inner_state_preserved_when_outer_has_no_conflict(self) -> None:
        """Inner's command keys are preserved when outer doesn't conflict."""

        class MyState(AgentState):
            inner_key: str
            outer_key: str

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

            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ExtendedModelResponse:
                response = handler(request)
                return ExtendedModelResponse(
                    model_response=response,
                    command=Command(update={"outer_key": "from_outer"}),
                )

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

            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ExtendedModelResponse:
                response = handler(request)
                return ExtendedModelResponse(
                    model_response=response,
                    command=Command(update={"inner_key": "from_inner"}),
                )

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

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

        # Both keys survive since there's no conflict
        messages = result["messages"]
        assert messages[-1].content == "Hello"