TestComposition Class — langchain Architecture

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_wrap_model_call_state_update.py lines 312–636

class TestComposition:
    """Test ExtendedModelResponse with composed middleware.

    Key semantics: Commands are collected inner-first, then outer.
    For non-reducer fields, later Commands overwrite (outer wins).
    For reducer fields (messages), all Commands are additive.
    """

    def test_outer_command_messages_added_alongside_model(self) -> None:
        """Outer middleware's command messages are added alongside model messages."""
        execution_order: list[str] = []

        class OuterMiddleware(AgentMiddleware):
            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ExtendedModelResponse:
                execution_order.append("outer-before")
                response = handler(request)
                execution_order.append("outer-after")
                return ExtendedModelResponse(
                    model_response=response,
                    command=Command(
                        update={"messages": [HumanMessage(content="Outer msg", id="outer-msg")]}
                    ),
                )

        class InnerMiddleware(AgentMiddleware):
            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ModelResponse:
                execution_order.append("inner-before")
                response = handler(request)
                execution_order.append("inner-after")
                return response

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

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

        # Execution order: outer wraps inner
        assert execution_order == [
            "outer-before",
            "inner-before",
            "inner-after",
            "outer-after",
        ]

        # Model messages + outer command messages (additive)
        messages = result["messages"]
        assert len(messages) == 3
        assert messages[0].content == "Hi"
        assert messages[1].content == "Composed"
        assert messages[2].content == "Outer msg"

    def test_inner_command_propagated_through_composition(self) -> None:
        """Inner middleware's ExtendedModelResponse command is propagated.

        When inner middleware returns ExtendedModelResponse, its command is
        captured before normalizing to ModelResponse at the composition boundary
        and collected into the final result.
        """

        class OuterMiddleware(AgentMiddleware):
            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ModelResponse:
                # Outer sees a ModelResponse from handler (inner's ExtendedModelResponse
                # was normalized at the composition boundary)
                response = handler(request)
                assert isinstance(response, ModelResponse)
                return response