TestCacheControlPreservation Class — langchain Architecture

Source Code

libs/langchain_v1/tests/unit_tests/agents/test_system_message.py lines 573–679

class TestCacheControlPreservation:
    """Test cache control metadata preservation in system messages."""

    def test_middleware_can_add_cache_control(self) -> None:
        """Test middleware adding cache control to system message."""

        def cache_control_middleware(
            request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]
        ) -> ModelResponse:
            """Add cache control to system message."""
            new_message = SystemMessage(
                content=[
                    {"type": "text", "text": "Base instructions"},
                    {
                        "type": "text",
                        "text": "Cached instructions",
                        "cache_control": {"type": "ephemeral"},
                    },
                ]
            )
            new_request = request.override(system_message=new_message)
            return handler(new_request)

        model = GenericFakeChatModel(messages=iter([AIMessage(content="response")]))
        request = ModelRequest(
            model=model,
            system_message=None,
            messages=[],
            tool_choice=None,
            tools=[],
            response_format=None,
            state=AgentState(messages=[]),
            runtime=Runtime(),
        )

        captured_request = None

        def mock_handler(req: ModelRequest) -> ModelResponse:
            nonlocal captured_request
            captured_request = req
            return ModelResponse(result=[AIMessage(content="response")])

        cache_control_middleware(request, mock_handler)

        assert captured_request is not None
        assert captured_request.system_message is not None
        assert isinstance(captured_request.system_message.content_blocks, list)
        assert captured_request.system_message.content_blocks[1].get("cache_control") == {
            "type": "ephemeral"
        }

    def test_cache_control_preserved_across_middleware(self) -> None:
        """Test that cache control is preserved when middleware modifies message."""

        def first_middleware_with_cache(
            request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]
        ) -> ModelResponse:
            """Set system message with cache control."""
            new_message = SystemMessage(
                content=[
                    {
                        "type": "text",
                        "text": "Cached content",
                        "cache_control": {"type": "ephemeral"},
                    }
                ]
            )
            new_request = request.override(system_message=new_message)
            return handler(new_request)

        def second_middleware_appends(
            request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]
        ) -> ModelResponse:
            """Append to system message while preserving cache control."""
            assert request.system_message is not None
            existing_content = request.system_message.content_blocks
            new_content = [*existing_content, TextContentBlock(type="text", text="Additional text")]

            new_message = SystemMessage(content_blocks=new_content)
            new_request = request.override(system_message=new_message)
            return handler(new_request)