tests/middleware_test.py

# -*- coding: utf-8 -*-
# pylint: disable=abstract-method
"""Unit tests for middleware system."""
from unittest.async_case import IsolatedAsyncioTestCase
from unittest.mock import AsyncMock, patch
from typing import Any, AsyncGenerator, Callable, Union

from utils import MockModel
from pydantic import BaseModel
from agentscope.event import AgentEvent
from agentscope.agent import Agent, ContextConfig
from agentscope.middleware import MiddlewareBase
from agentscope.model import ChatResponse
from agentscope.message import (
    TextBlock,
    HintBlock,
    UserMsg,
    SystemMsg,
    Msg,
    ToolCallBlock,
)
from agentscope.tool import Toolkit, ToolBase, ToolChunk
from agentscope.permission import (
    PermissionContext,
    PermissionDecision,
    PermissionBehavior,
)


class TestMiddleware(IsolatedAsyncioTestCase):
    """Test cases for middleware system."""

    async def asyncSetUp(self) -> None:
        """Set up test fixtures."""
        self.mock_model = MockModel()
        self.toolkit = Toolkit()
        self.execution_log = []

    async def test_on_reply_middleware_pre_post_yield(self) -> None:
        """Test on_reply middleware pre, post and yield positions."""

        class ReplyMiddleware(MiddlewareBase):
            """Middleware for testing on_reply hook."""

            def __init__(self, log: list, name: str) -> None:
                """Initialize the reply middleware.

                Args:
                    log: The execution log list.
                    name: The middleware name.
                """
                self.log = log
                self.name = name

            async def on_reply(
                self,
                agent: Agent,
                input_kwargs: dict,
                next_handler: Callable[[], AsyncGenerator],
            ) -> AsyncGenerator:
                """The on_reply middleware logic."""
                self.log.append(f"{self.name}_pre")
                async for item in next_handler():
                    if isinstance(item, AgentEvent):
                        self.log.append(f"{self.name}_{item.type}")
                    elif isinstance(item, Msg):
                        self.log.append(f"{self.name}_msg")
                    yield item
                self.log.append(f"{self.name}_post")

        middleware1 = ReplyMiddleware(self.execution_log, "mw1")
        middleware2 = ReplyMiddleware(self.execution_log, "mw2")

        self.mock_model.set_responses(
            [
                ChatResponse(
                    content=[TextBlock(text="test response")],
                    is_last=True,
                ),
            ],
        )

        agent = Agent(
            name="test_agent",
            system_prompt="test prompt",
            model=self.mock_model,
            toolkit=self.toolkit,
            middlewares=[middleware1, middleware2],
        )

        await agent.reply(UserMsg("user", "test message"))

        # Verify execution order
        expected = [
            "mw1_pre",
            "mw2_pre",
            "mw2_REPLY_START",
            "mw1_REPLY_START",
            "mw2_MODEL_CALL_START",
            "mw1_MODEL_CALL_START",
            "mw2_TEXT_BLOCK_START",
            "mw1_TEXT_BLOCK_START",
            "mw2_TEXT_BLOCK_DELTA",
            "mw1_TEXT_BLOCK_DELTA",
            "mw2_TEXT_BLOCK_END",
            "mw1_TEXT_BLOCK_END",
            "mw2_MODEL_CALL_END",
            "mw1_MODEL_CALL_END",
            "mw2_REPLY_END",
            "mw1_REPLY_END",
            "mw2_msg",
            "mw1_msg",
            "mw2_post",
            "mw1_post",
        ]
        self.assertListEqual(self.execution_log, expected)

    async def test_on_reasoning_middleware_pre_yield(self) -> None:
        """Test on_reasoning middleware pre and yield positions."""

        class ReasoningMiddleware(MiddlewareBase):
            """Middleware for testing on_reasoning hook."""

            def __init__(self, log: list, name: str) -> None:
                """Initialize the reasoning middleware.

                Args:
                    log: The execution log list.
                    name: The middleware name.
                """
                self.log = log
                self.name = name

            async def on_reasoning(
                self,
                agent: Agent,
                input_kwargs: dict,
                next_handler: Callable[[], AsyncGenerator],
            ) -> AsyncGenerator:
                """The on_reasoning middleware logic."""
                self.log.append(f"{self.name}_pre")
                async for item in next_handler():
                    if isinstance(item, AgentEvent):
                        self.log.append(f"{self.name}_{item.type}")
                    elif isinstance(item, Msg):
                        self.log.append(f"{self.name}_msg")
                    yield item
                self.log.append(f"{self.name}_post")

        middleware1 = ReasoningMiddleware(self.execution_log, "mw1")
        middleware2 = ReasoningMiddleware(self.execution_log, "mw2")

        self.mock_model.set_responses(
            [
                ChatResponse(
                    content=[TextBlock(text="test response")],
                    is_last=True,
                ),
            ],
        )

        agent = Agent(
            name="test_agent",
            system_prompt="test prompt",
            model=self.mock_model,
            toolkit=self.toolkit,
            middlewares=[middleware1, middleware2],
        )

        await agent.reply(UserMsg("user", "test message"))

        # Verify execution order
        expected = [
            "mw1_pre",
            "mw2_pre",
            "mw2_MODEL_CALL_START",
            "mw1_MODEL_CALL_START",
            "mw2_TEXT_BLOCK_START",
            "mw1_TEXT_BLOCK_START",
            "mw2_TEXT_BLOCK_DELTA",
            "mw1_TEXT_BLOCK_DELTA",
            "mw2_TEXT_BLOCK_END",
            "mw1_TEXT_BLOCK_END",
            "mw2_MODEL_CALL_END",
            "mw1_MODEL_CALL_END",
            "mw2_msg",
            "mw1_msg",
        ]
        self.assertListEqual(self.execution_log, expected)

    async def test_on_model_call_middleware_non_streaming(self) -> None:
        """Test on_model_call middleware for non-streaming model."""

        class ModelCallMiddleware(MiddlewareBase):
            """Middleware for testing on_model_call hook."""

            def __init__(self, log: list, name: str) -> None:
                """Initialize the model call middleware.

                Args:
                    log: The execution log list.
                    name: The middleware name.
                """
                self.log = log
                self.name = name

            async def on_model_call(
                self,
                agent: Agent,
                input_kwargs: dict,
                next_handler: Callable,
            ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]:
                """The on_model_call middleware logic."""
                self.log.append(f"{self.name}_pre")
                result = await next_handler()
                self.log.append(f"{self.name}_post")
                return result

        middleware1 = ModelCallMiddleware(self.execution_log, "mw1")
        middleware2 = ModelCallMiddleware(self.execution_log, "mw2")

        self.mock_model.set_responses(
            [
                ChatResponse(
                    content=[TextBlock(text="test response")],
                    is_last=True,
                ),
            ],
        )

        agent = Agent(
            name="test_agent",
            system_prompt="test prompt",
            model=self.mock_model,
            toolkit=self.toolkit,
            middlewares=[middleware1, middleware2],
        )

        await agent.reply(UserMsg("user", "test message"))

        # Verify execution order: mw1_pre -> mw2_pre -> mw2_post -> mw1_post
        expected = ["mw1_pre", "mw2_pre", "mw2_post", "mw1_post"]
        self.assertListEqual(self.execution_log, expected)

    async def test_on_model_call_middleware_streaming(self) -> None:
        """Test on_model_call middleware for streaming model."""

        class ModelCallMiddleware(MiddlewareBase):
            """Middleware for testing on_model_call hook with streaming."""

            def __init__(self, log: list, name: str) -> None:
                """Initialize the model call middleware.

                Args:
                    log: The execution log list.
                    name: The middleware name.
                """
                self.log = log
                self.name = name

            async def on_model_call(
                self,
                agent: Agent,
                input_kwargs: dict,
                next_handler: Callable,
            ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]:
                """The on_model_call middleware logic for streaming."""
                self.log.append(f"{self.name}_pre")
                result = await next_handler()

                async def wrapped_generator() -> AsyncGenerator[
                    ChatResponse,
                    None,
                ]:
                    """Wrap the generator to log yields."""
                    async for chunk in result:
                        self.log.append(f"{self.name}_chunk")
                        yield chunk
                    self.log.append(f"{self.name}_post")

                return wrapped_generator()

        middleware1 = ModelCallMiddleware(self.execution_log, "mw1")
        middleware2 = ModelCallMiddleware(self.execution_log, "mw2")

        self.mock_model.set_responses(
            [
                [
                    ChatResponse(
                        content=[TextBlock(text="chunk1")],
                        is_last=False,
                    ),
                    ChatResponse(
                        content=[TextBlock(text="chunk2")],
                        is_last=False,
                    ),
                    ChatResponse(
                        content=[TextBlock(text="chunk3")],
                        is_last=True,
                    ),
                ],
            ],
        )

        agent = Agent(
            name="test_agent",
            system_prompt="test prompt",
            model=self.mock_model,
            toolkit=self.toolkit,
            middlewares=[middleware1, middleware2],
        )

        await agent.reply(UserMsg("user", "test message"))

        # Verify execution order
        expected = [
            "mw1_pre",
            "mw2_pre",
            "mw2_chunk",
            "mw1_chunk",
            "mw2_chunk",
            "mw1_chunk",
            "mw2_chunk",
            "mw1_chunk",
            "mw2_post",
            "mw1_post",
        ]
        self.assertListEqual(self.execution_log, expected)

    async def test_on_system_prompt_middleware(self) -> None:
        """Test on_system_prompt middleware (transformer pattern)."""

        class SystemPromptMiddleware(MiddlewareBase):
            """Middleware for testing on_system_prompt hook."""

            def __init__(self, log: list, name: str, suffix: str) -> None:
                """Initialize the system prompt middleware.

                Args:
                    log: The execution log list.
                    name: The middleware name.
                    suffix: The suffix to append to the prompt.
                """
                self.log = log
                self.name = name
                self.suffix = suffix

            async def on_system_prompt(
                self,
                agent: Agent,
                current_prompt: str,
            ) -> str:
                """The on_system_prompt middleware logic."""
                self.log.append(f"{self.name}_executed")
                return f"{current_prompt} {self.suffix}"

        middleware1 = SystemPromptMiddleware(
            self.execution_log,
            "mw1",
            "[MW1]",
        )
        middleware2 = SystemPromptMiddleware(
            self.execution_log,
            "mw2",
            "[MW2]",
        )

        self.mock_model.set_responses(
            [
                ChatResponse(
                    content=[TextBlock(text="test response")],
                    is_last=True,
                ),
            ],
        )

        agent = Agent(
            name="test_agent",
            system_prompt="test prompt",
            model=self.mock_model,
            toolkit=self.toolkit,
            middlewares=[middleware1, middleware2],
        )

        await agent.reply(UserMsg("user", "test message"))

        # Verify execution order: mw1 -> mw2 (sequential transformer pattern)
        # Note: system_prompt is called twice (once for initial setup, once
        # during reasoning)
        expected = [
            "mw1_executed",
            "mw2_executed",  # First call
            "mw1_executed",
            "mw2_executed",  # Second call during reasoning
        ]
        self.assertListEqual(self.execution_log, expected)

    async def test_multiple_middleware_types(self) -> None:
        """Test multiple middleware types working together."""

        class MultiMiddleware(MiddlewareBase):
            """Middleware implementing multiple hooks."""

            def __init__(self, log: list, name: str) -> None:
                """Initialize the multi middleware.

                Args:
                    log: The execution log list.
                    name: The middleware name.
                """
                self.log = log
                self.name = name

            async def on_reply(
                self,
                agent: Agent,
                input_kwargs: dict,
                next_handler: Callable[[], AsyncGenerator],
            ) -> AsyncGenerator:
                """The on_reply middleware logic."""
                self.log.append(f"{self.name}_reply_pre")
                async for item in next_handler():
                    yield item
                self.log.append(f"{self.name}_reply_post")

            async def on_reasoning(
                self,
                agent: Agent,
                input_kwargs: dict,
                next_handler: Callable[[], AsyncGenerator],
            ) -> AsyncGenerator:
                """The on_reasoning middleware logic."""
                self.log.append(f"{self.name}_reasoning_pre")
                async for item in next_handler():
                    yield item
                self.log.append(f"{self.name}_reasoning_post")

            async def on_model_call(
                self,
                agent: Agent,
                input_kwargs: dict,
                next_handler: Callable,
            ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]:
                """The on_model_call middleware logic."""
                self.log.append(f"{self.name}_model_call_pre")
                result = await next_handler()
                self.log.append(f"{self.name}_model_call_post")
                return result

            async def on_system_prompt(
                self,
                agent: Agent,
                current_prompt: str,
            ) -> str:
                """The on_system_prompt middleware logic."""
                self.log.append(f"{self.name}_system_prompt")
                return current_prompt

        middleware = MultiMiddleware(self.execution_log, "multi")

        self.mock_model.set_responses(
            [
                ChatResponse(
                    content=[TextBlock(text="test response")],
                    is_last=True,
                ),
            ],
        )

        agent = Agent(
            name="test_agent",
            system_prompt="test prompt",
            model=self.mock_model,
            toolkit=self.toolkit,
            middlewares=[middleware],
        )

        await agent.reply(UserMsg("user", "test message"))

        # Verify all middleware hooks were called
        expected = [
            "multi_reply_pre",
            "multi_system_prompt",
            "multi_reasoning_pre",
            "multi_system_prompt",
            "multi_model_call_pre",
            "multi_model_call_post",
            "multi_reply_post",
        ]
        self.assertListEqual(self.execution_log, expected)

    async def test_on_reply_middleware_modify_input(self) -> None:
        """Test that on_reply middleware can modify msgs input."""

        class ModifyMsgsMiddleware(MiddlewareBase):
            """Middleware that modifies the msgs input."""

            async def on_reply(
                self,
                agent: Agent,
                input_kwargs: dict,
                next_handler: Callable[..., AsyncGenerator],
            ) -> AsyncGenerator:
                """Modify inputs before passing to next handler."""
                # Modify the message content
                inputs = input_kwargs["inputs"]
                if isinstance(inputs, Msg):
                    modified_msg = UserMsg(
                        name=inputs.name,
                        content="MODIFIED: " + inputs.get_text_content(),
                    )
                    async for item in next_handler(inputs=modified_msg):
                        yield item
                else:
                    async for item in next_handler(**input_kwargs):
                        yield item

        middleware = ModifyMsgsMiddleware()

        # Track what message the model receives
        received_messages = []

        class TrackingModel(MockModel):
            """Model that tracks received messages."""

            async def _call_api(
                self,
                *args: Any,
                **kwargs: Any,
            ) -> ChatResponse:
                """Track the messages and return mock response."""
                messages = kwargs.get("messages", [])
                received_messages.extend(messages)
                return await super()._call_api(*args, **kwargs)

        tracking_model = TrackingModel()
        tracking_model.set_responses(
            [
                ChatResponse(
                    content=[TextBlock(text="response")],
                    is_last=True,
                ),
            ],
        )

        agent_instance = Agent(
            name="test_agent",
            system_prompt="test prompt",
            model=tracking_model,
            toolkit=self.toolkit,
            middlewares=[middleware],
        )

        await agent_instance.reply(UserMsg("user", "original message"))

        # Verify the model received the modified message
        user_messages = [m for m in received_messages if m.role == "user"]
        self.assertTrue(len(user_messages) > 0)
        self.assertIn(
            "MODIFIED: original message",
            user_messages[-1].get_text_content(),
        )

    async def test_on_reasoning_middleware_modify_input(self) -> None:
        """Test that on_reasoning middleware can modify tool_choice input."""

        class ModifyToolChoiceMiddleware(MiddlewareBase):
            """Middleware that modifies the tool_choice input."""

            async def on_reasoning(
                self,
                agent: Agent,
                input_kwargs: dict,
                next_handler: Callable[..., AsyncGenerator],
            ) -> AsyncGenerator:
                """Force tool_choice to 'none' to prevent tool calls."""
                # Override tool_choice to 'none'
                async for item in next_handler(tool_choice="none"):
                    yield item

        middleware = ModifyToolChoiceMiddleware()

        # Track what tool_choice the model receives
        received_tool_choices = []

        class TrackingModel(MockModel):
            """Model that tracks received tool_choice."""

            async def _call_api(
                self,
                *args: Any,
                **kwargs: Any,
            ) -> ChatResponse:
                """Track the tool_choice and return mock response."""
                tool_choice = kwargs.get("tool_choice")
                received_tool_choices.append(tool_choice)
                return await super()._call_api(*args, **kwargs)

        tracking_model = TrackingModel()
        tracking_model.set_responses(
            [
                ChatResponse(
                    content=[TextBlock(text="response without tools")],
                    is_last=True,
                ),
            ],
        )

        agent_instance = Agent(
            name="test_agent",
            system_prompt="test prompt",
            model=tracking_model,
            toolkit=self.toolkit,
            middlewares=[middleware],
        )

        await agent_instance.reply(UserMsg("user", "test message"))

        # Verify the model received tool_choice='none'
        self.assertIn("none", received_tool_choices)

    async def test_on_acting_middleware_intercepts_tool_execution(
        self,
    ) -> None:
        """Test that on_acting middleware intercepts raw tool execution.

        After the refactor, ``on_acting`` wraps only ``_acting_impl``
        (i.e. ``toolkit.call_tool``).  Permission checking and context
        writes are handled by ``_execute_tool_call`` *outside* the hook.
        This test verifies that the middleware can observe and modify the
        ``tool_call`` passed to the actual tool function.
        """

        # ------------------------------------------------------------------ #
        # A minimal tool that records the raw input it receives.              #
        # ------------------------------------------------------------------ #
        received_inputs: list[str] = []

        class _EchoParams(BaseModel):
            value: str

        class EchoTool(ToolBase):
            """Tool that echoes its input and records it."""

            name: str = "echo"
            description: str = "Echo the value."
            input_schema: dict = _EchoParams.model_json_schema()
            is_concurrency_safe: bool = True
            is_read_only: bool = True
            is_state_injected: bool = False
            is_external_tool: bool = False
            is_mcp: bool = False
            mcp_name: str | None = None

            async def check_permissions(
                self,
                tool_input: dict[str, Any],
                context: PermissionContext,
            ) -> PermissionDecision:
                """Always allow."""
                return PermissionDecision(
                    behavior=PermissionBehavior.ALLOW,
                    message="allowed",
                )

            async def __call__(
                self,
                value: str,
            ) -> ToolChunk:
                """Record the value and return it."""
                received_inputs.append(value)
                return ToolChunk(
                    content=[TextBlock(text=f"echo:{value}")],
                )

        toolkit_with_tool = Toolkit(tools=[EchoTool()])

        # ------------------------------------------------------------------ #
        # Middleware that renames the tool_call.input before forwarding.      #
        # ------------------------------------------------------------------ #
        intercepted_tool_calls: list[str] = []

        class ObserveActingMiddleware(MiddlewareBase):
            """Middleware that records the tool_call seen at acting level."""

            async def on_acting(
                self,
                agent: Agent,
                input_kwargs: dict,
                next_handler: Callable[..., AsyncGenerator],
            ) -> AsyncGenerator:
                """Record the tool_call and forward to next handler."""
                tool_call = input_kwargs["tool_call"]
                intercepted_tool_calls.append(tool_call.input)
                # Modify the input before execution
                import json

                modified = ToolCallBlock(
                    id=tool_call.id,
                    name=tool_call.name,
                    input=json.dumps({"value": "MODIFIED"}),
                    state=tool_call.state,
                )
                async for chunk in next_handler(tool_call=modified):
                    yield chunk

        middleware = ObserveActingMiddleware()
        self.mock_model.set_responses(
            [
                ChatResponse(
                    content=[TextBlock(text="done")],
                    is_last=True,
                ),
            ],
        )

        agent_instance = Agent(
            name="test_agent",
            system_prompt="test prompt",
            model=self.mock_model,
            toolkit=toolkit_with_tool,
            middlewares=[middleware],
        )

        # Call _execute_tool_call with a valid tool call.
        tool_call = ToolCallBlock(
            id="call_1",
            name="echo",
            input='{"value": "ORIGINAL"}',
        )
        events = []
        # pylint: disable=protected-access
        async for evt in agent_instance._execute_tool_call(tool_call):
            events.append(evt)

        # Middleware intercepted the tool call at execution level
        self.assertEqual(len(intercepted_tool_calls), 1)
        self.assertIn("ORIGINAL", intercepted_tool_calls[0])

        # The tool actually received the MODIFIED value
        self.assertEqual(len(received_inputs), 1)
        self.assertEqual(received_inputs[0], "MODIFIED")

    async def test_on_model_call_middleware_modify_input(self) -> None:
        """Test that on_model_call middleware can modify messages and model."""

        class ModifyMessagesMiddleware(MiddlewareBase):
            """Middleware that modifies messages input."""

            async def on_model_call(
                self,
                agent: Agent,
                input_kwargs: dict,
                next_handler: Callable[..., Any],
            ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]:
                """Prepend a system message to the messages list."""
                messages = input_kwargs["messages"]
                modified_messages = [
                    SystemMsg(
                        name="system",
                        content="INJECTED SYSTEM MESSAGE",
                    ),
                ] + messages

                # Pass modified messages to next handler
                return await next_handler(
                    current_model=input_kwargs["current_model"],
                    messages=modified_messages,
                    tools=input_kwargs["tools"],
                    tool_choice=input_kwargs["tool_choice"],
                )

        middleware = ModifyMessagesMiddleware()

        # Track what messages the model receives
        received_messages = []

        class TrackingModel(MockModel):
            """Model that tracks received messages."""

            async def _call_api(
                self,
                *args: Any,
                **kwargs: Any,
            ) -> ChatResponse:
                """Track the messages and return mock response."""
                messages = kwargs.get("messages", [])
                received_messages.extend(messages)
                return await super()._call_api(*args, **kwargs)

        tracking_model = TrackingModel()
        tracking_model.set_responses(
            [
                ChatResponse(
                    content=[TextBlock(text="response")],
                    is_last=True,
                ),
            ],
        )

        agent_instance = Agent(
            name="test_agent",
            system_prompt="test prompt",
            model=tracking_model,
            toolkit=self.toolkit,
            middlewares=[middleware],
        )

        await agent_instance.reply(UserMsg("user", "test message"))

        # Verify the injected system message is present
        system_messages = [m for m in received_messages if m.role == "system"]
        self.assertTrue(
            any(
                "INJECTED SYSTEM MESSAGE" in m.get_text_content()
                for m in system_messages
            ),
        )

    async def test_on_compress_context_middleware(self) -> None:
        """Test on_compress_context middleware follows the onion chain pattern.

        Verifies that:
        - Multiple middlewares are chained in onion order (mw1 wraps mw2).
        - ``input_kwargs`` carries the correct ``context_config`` and
          ``instructions``.
        - The ``next_handler`` ultimately calls ``_compress_context_impl``.
        - A middleware can short-circuit and skip the actual implementation.
        """
        seen_instructions = []

        # ------------------------------------------------------------------ #
        # Middleware that records pre/post and forwards to next_handler.      #
        # ------------------------------------------------------------------ #
        class CompressContextMiddleware(MiddlewareBase):
            """Middleware for testing on_compress_context hook."""

            def __init__(self, log: list, name: str) -> None:
                """Initialize the compress context middleware.

                Args:
                    log (`list`):
                        The execution log list.
                    name (`str`):
                        The middleware name.
                """
                self.log = log
                self.name = name

            async def on_compress_context(
                self,
                agent: Agent,
                input_kwargs: dict,
                next_handler: Callable[..., Any],
            ) -> None:
                """Forward to next handler, recording pre and post."""
                self.log.append(f"{self.name}_pre")
                seen_instructions.append(input_kwargs.get("instructions"))
                await next_handler(**input_kwargs)
                self.log.append(f"{self.name}_post")

        middleware1 = CompressContextMiddleware(self.execution_log, "mw1")
        middleware2 = CompressContextMiddleware(self.execution_log, "mw2")

        self.mock_model.set_responses(
            [
                ChatResponse(
                    content=[TextBlock(text="test response")],
                    is_last=True,
                ),
            ],
        )

        context_config = ContextConfig(trigger_ratio=0.8, reserve_ratio=0.1)
        agent = Agent(
            name="test_agent",
            system_prompt="test prompt",
            model=self.mock_model,
            toolkit=self.toolkit,
            middlewares=[middleware1, middleware2],
            context_config=context_config,
        )

        # Patch _compress_context_impl to avoid real token counting.
        instructions = HintBlock(
            hint="Keep user requirements while compressing.",
            source="user",
        )
        with patch.object(
            agent,
            "_compress_context_impl",
            new_callable=AsyncMock,
        ) as mock_impl:
            await agent.compress_context(
                context_config=context_config,
                instructions=instructions,
            )

            # _compress_context_impl must have been called exactly once.
            mock_impl.assert_awaited_once_with(
                context_config=context_config,
                instructions=instructions,
            )

        # Verify onion execution order: mw1_pre -> mw2_pre -> mw2_post ->
        # mw1_post
        expected = ["mw1_pre", "mw2_pre", "mw2_post", "mw1_post"]
        self.assertListEqual(self.execution_log, expected)
        self.assertListEqual(seen_instructions, [instructions, instructions])

    async def test_on_compress_context_middleware_modify_instructions(
        self,
    ) -> None:
        """Test that middleware can replace compress_context instructions."""

        class ReplaceInstructionsMiddleware(MiddlewareBase):
            """Middleware that replaces the compression instructions."""

            def __init__(self, replacement: HintBlock) -> None:
                """Initialize the middleware with replacement instructions."""
                self.replacement = replacement

            async def on_compress_context(
                self,
                agent: Agent,
                input_kwargs: dict,
                next_handler: Callable[..., Any],
            ) -> None:
                """Replace instructions before forwarding."""
                input_kwargs["instructions"] = self.replacement
                await next_handler(**input_kwargs)

        original = HintBlock(
            hint="Keep all requirements.",
            source="user",
        )
        replacement = HintBlock(
            hint="Keep only unresolved requirements.",
            source="middleware",
        )
        context_config = ContextConfig(trigger_ratio=0.8, reserve_ratio=0.1)
        agent = Agent(
            name="test_agent",
            system_prompt="test prompt",
            model=self.mock_model,
            toolkit=self.toolkit,
            middlewares=[ReplaceInstructionsMiddleware(replacement)],
            context_config=context_config,
        )

        with patch.object(
            agent,
            "_compress_context_impl",
            new_callable=AsyncMock,
        ) as mock_impl:
            await agent.compress_context(
                context_config=context_config,
                instructions=original,
            )

            mock_impl.assert_awaited_once_with(
                context_config=context_config,
                instructions=replacement,
            )

    async def test_on_compress_context_middleware_short_circuit(
        self,
    ) -> None:
        """Test that a middleware can skip _compress_context_impl entirely."""

        class SkipCompressMiddleware(MiddlewareBase):
            """Middleware that skips the actual compress_context call."""

            def __init__(self, log: list) -> None:
                """Initialize the skip compress middleware.

                Args:
                    log (`list`):
                        The execution log list.
                """
                self.log = log

            async def on_compress_context(
                self,
                agent: Agent,
                input_kwargs: dict,
                next_handler: Callable[..., Any],
            ) -> None:
                """Record the call and skip forwarding to next_handler."""
                self.log.append("skipped")
                # Intentionally NOT calling next_handler.

        middleware = SkipCompressMiddleware(self.execution_log)

        self.mock_model.set_responses(
            [
                ChatResponse(
                    content=[TextBlock(text="test response")],
                    is_last=True,
                ),
            ],
        )

        agent = Agent(
            name="test_agent",
            system_prompt="test prompt",
            model=self.mock_model,
            toolkit=self.toolkit,
            middlewares=[middleware],
        )

        with patch.object(
            agent,
            "_compress_context_impl",
            new_callable=AsyncMock,
        ) as mock_impl:
            await agent.compress_context()

            # _compress_context_impl should NOT have been called.
            mock_impl.assert_not_awaited()

        self.assertListEqual(self.execution_log, ["skipped"])

    async def test_on_compress_context_no_middleware(self) -> None:
        """Test that compress_context calls _compress_context_impl directly
        when no middleware is registered."""

        self.mock_model.set_responses(
            [
                ChatResponse(
                    content=[TextBlock(text="test response")],
                    is_last=True,
                ),
            ],
        )

        context_config = ContextConfig(trigger_ratio=0.8, reserve_ratio=0.1)
        agent = Agent(
            name="test_agent",
            system_prompt="test prompt",
            model=self.mock_model,
            toolkit=self.toolkit,
            # No middlewares registered.
            context_config=context_config,
        )

        with patch.object(
            agent,
            "_compress_context_impl",
            new_callable=AsyncMock,
        ) as mock_impl:
            await agent.compress_context(context_config=context_config)
            mock_impl.assert_awaited_once_with(
                context_config=context_config,
                instructions=None,
            )

    async def asyncTearDown(self) -> None:
        """Clean up test fixtures."""
        self.execution_log.clear()