agentscope-ai--agentscope/tests/middleware_budget_test.py

# -*- coding: utf-8 -*-
"""Unit tests for BudgetControlMiddleware."""
from typing import Any
from unittest.async_case import IsolatedAsyncioTestCase

from utils import MockModel
from agentscope.agent import Agent
from agentscope.message import UserMsg, TextBlock, ToolCallBlock, HintBlock
from agentscope.middleware import ReplyBudgetControlMiddleware
from agentscope.model import ChatResponse, ChatUsage
from agentscope.permission import (
    PermissionBehavior,
    PermissionContext,
    PermissionDecision,
)
from agentscope.event import UserConfirmResultEvent, ConfirmResult
from agentscope.tool import ToolBase, Toolkit, ToolChunk


def _response(
    text: str,
    input_tokens: int,
    output_tokens: int,
) -> ChatResponse:
    """Build a non-streaming ChatResponse with usage."""
    return ChatResponse(
        content=[TextBlock(text=text)],
        is_last=True,
        usage=ChatUsage(
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            time=0.0,
        ),
    )


class DummyTool(ToolBase):
    """Minimal tool that always allows and returns a fixed result."""

    name: str = "dummy"
    description: str = "A dummy tool for testing"
    input_schema: dict[str, Any] = {"type": "object", "properties": {}}
    is_concurrency_safe: bool = True
    is_read_only: bool = True
    is_external_tool: bool = False
    is_mcp: bool = False

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

    async def __call__(self, **kwargs: Any) -> ToolChunk:
        """Return a fixed result."""
        return ToolChunk(content=[TextBlock(text="ok")])


class ConfirmRequiredTool(ToolBase):
    """Minimal tool that always requires user confirmation before running."""

    name: str = "confirm_required"
    description: str = "A tool that requires user confirmation"
    input_schema: dict[str, Any] = {"type": "object", "properties": {}}
    is_concurrency_safe: bool = False
    is_read_only: bool = False
    is_external_tool: bool = False
    is_mcp: bool = False

    async def check_permissions(
        self,
        tool_input: dict[str, Any],
        context: PermissionContext,
    ) -> PermissionDecision:
        """Always require user confirmation."""
        return PermissionDecision(
            behavior=PermissionBehavior.ASK,
            decision_reason="Confirmation required",
            message="Confirmation required",
        )

    async def __call__(self, **kwargs: Any) -> ToolChunk:
        """Return a fixed result."""
        return ToolChunk(content=[TextBlock(text="confirmed result")])


def _has_hint_block(msg: Any, hint_message: str) -> bool:
    """Return True if *msg* contains a HintBlock with *hint_message*."""
    content = getattr(msg, "content", None)
    if not isinstance(content, list):
        return False
    return any(
        isinstance(b, HintBlock) and hint_message in b.hint for b in content
    )


class TestBudgetControlMiddleware(IsolatedAsyncioTestCase):
    """Test cases for BudgetControlMiddleware."""

    async def asyncSetUp(self) -> None:
        """Set up shared fixtures."""
        self.toolkit = Toolkit()

    async def test_under_budget_no_hint_injected(self) -> None:
        """When token usage stays below the budget, no hint is injected."""
        model = MockModel()
        model.set_responses(
            [_response("done", input_tokens=10, output_tokens=5)],
        )

        middleware = ReplyBudgetControlMiddleware(token_budget=1000)
        agent = Agent(
            name="test_agent",
            system_prompt="you are helpful",
            model=model,
            toolkit=self.toolkit,
            middlewares=[middleware],
        )

        context_before = len(agent.state.context)
        await agent.reply(UserMsg("user", "hello"))

        # No HintBlock should have been added to context
        hint_msgs = [
            m
            for m in agent.state.context[context_before:]
            if _has_hint_block(m, middleware.hint_message)
        ]
        self.assertEqual(len(hint_msgs), 0)

    async def test_budget_exceeded_injects_hint(self) -> None:
        """When the budget is exceeded, the hint block is injected.

        Uses token_budget=0 so the budget condition fires on the very first
        reasoning call (0 used >= 0 max).
        """
        model = MockModel()
        model.set_responses(
            [_response("wrap up", input_tokens=10, output_tokens=5)],
        )

        middleware = ReplyBudgetControlMiddleware(token_budget=0)
        agent = Agent(
            name="test_agent",
            system_prompt="you are helpful",
            model=model,
            toolkit=self.toolkit,
            middlewares=[middleware],
        )

        context_before = len(agent.state.context)
        await agent.reply(UserMsg("user", "hello"))

        hint_msgs = [
            m
            for m in agent.state.context[context_before:]
            if _has_hint_block(m, middleware.hint_message)
        ]
        self.assertGreater(len(hint_msgs), 0)

    async def test_budget_exceeded_forces_tool_choice_none(self) -> None:
        """When budget is exceeded, tool_choice forwarded to model is none.

        Uses token_budget=0 so the override fires on the first reasoning call.
        """
        received_tool_choices: list = []

        class TrackingModel(MockModel):
            """Model that records tool_choice on every call."""

            async def _call_api(
                self,
                *args: Any,
                **kwargs: Any,
            ) -> ChatResponse:
                """Record tool_choice and delegate to mock."""
                received_tool_choices.append(kwargs.get("tool_choice"))
                return await super()._call_api(*args, **kwargs)

        model = TrackingModel()
        model.set_responses(
            [_response("wrap up", input_tokens=10, output_tokens=5)],
        )

        middleware = ReplyBudgetControlMiddleware(token_budget=0)
        agent = Agent(
            name="test_agent",
            system_prompt="you are helpful",
            model=model,
            toolkit=self.toolkit,
            middlewares=[middleware],
        )

        await agent.reply(UserMsg("user", "hello"))

        # At least one call must have received tool_choice with mode="none"
        self.assertTrue(
            any(
                getattr(tc, "mode", None) == "none"
                for tc in received_tool_choices
                if tc is not None
            ),
        )

    async def test_token_accumulation_across_steps(self) -> None:
        """Tokens accumulate across steps and trigger enforcement correctly.

        Step 1: tool call costs 200+100=300 tokens (token_budget=300 so
        step 2 sees used >= max and injects the hint).
        """
        toolkit = Toolkit(tools=[DummyTool()])

        model = MockModel()
        model.set_responses(
            [
                [
                    ChatResponse(
                        content=[
                            ToolCallBlock(
                                id="tc_1",
                                name="dummy",
                                input="{}",
                            ),
                        ],
                        is_last=True,
                        usage=ChatUsage(
                            input_tokens=200,
                            output_tokens=100,
                            time=0.0,
                        ),
                    ),
                ],
                [
                    ChatResponse(
                        content=[TextBlock(text="done")],
                        is_last=True,
                        usage=ChatUsage(
                            input_tokens=150,
                            output_tokens=50,
                            time=0.0,
                        ),
                    ),
                ],
            ],
        )

        middleware = ReplyBudgetControlMiddleware(token_budget=300)
        agent = Agent(
            name="test_agent",
            system_prompt="you are helpful",
            model=model,
            toolkit=toolkit,
            middlewares=[middleware],
        )

        context_before = len(agent.state.context)
        await agent.reply(UserMsg("user", "hello"))

        hint_msgs = [
            m
            for m in agent.state.context[context_before:]
            if _has_hint_block(m, middleware.hint_message)
        ]
        self.assertGreater(len(hint_msgs), 0)

    async def test_weighted_token_calculation(self) -> None:
        """output_token_weight scales output tokens in the budget calculation.

        With input_token_weight=1, output_token_weight=3, token_budget=200:
        - Step 1: 50 input * 1 + 50 output * 3 = 200 → budget hit exactly
        - Step 2: hint should be injected before the model call
        """
        toolkit = Toolkit(tools=[DummyTool()])

        model = MockModel()
        model.set_responses(
            [
                [
                    ChatResponse(
                        content=[
                            ToolCallBlock(
                                id="tc_1",
                                name="dummy",
                                input="{}",
                            ),
                        ],
                        is_last=True,
                        usage=ChatUsage(
                            input_tokens=50,
                            output_tokens=50,
                            time=0.0,
                        ),
                    ),
                ],
                [
                    ChatResponse(
                        content=[TextBlock(text="done")],
                        is_last=True,
                        usage=ChatUsage(
                            input_tokens=30,
                            output_tokens=10,
                            time=0.0,
                        ),
                    ),
                ],
            ],
        )

        # 50*1 + 50*3 = 200 == token_budget → step 2 triggers enforcement
        middleware = ReplyBudgetControlMiddleware(
            token_budget=200,
            input_token_weight=1,
            output_token_weight=3,
        )
        agent = Agent(
            name="test_agent",
            system_prompt="you are helpful",
            model=model,
            toolkit=toolkit,
            middlewares=[middleware],
        )

        context_before = len(agent.state.context)
        await agent.reply(UserMsg("user", "hello"))

        hint_msgs = [
            m
            for m in agent.state.context[context_before:]
            if _has_hint_block(m, middleware.hint_message)
        ]
        self.assertGreater(len(hint_msgs), 0)

    async def test_state_cleanup_after_reply(self) -> None:
        """middle_context entry for the reply is removed after reply ends."""
        model = MockModel()
        model.set_responses(
            [_response("done", input_tokens=10, output_tokens=5)],
        )

        middleware = ReplyBudgetControlMiddleware(token_budget=1000)
        agent = Agent(
            name="test_agent",
            system_prompt="you are helpful",
            model=model,
            toolkit=self.toolkit,
            middlewares=[middleware],
        )

        await agent.reply(UserMsg("user", "hello"))

        middleware_key = await middleware.get_middleware_key()
        bucket = agent.state.middle_context.get(middleware_key, {})
        # All per-reply entries must have been cleaned up
        self.assertEqual(len(bucket), 0)

    async def test_token_accumulation_persists_across_hitl(self) -> None:
        """Token accumulation in middle_context persists across HITL boundary.

        Scenario:
        - token_budget=300, both weights default to 1.
        - First reply_stream call: model call costs 200 input + 100 output
          = 300 tokens, then pauses at REQUIRE_USER_CONFIRM (no ReplyEndEvent
          is emitted). The 300-token count is stored in middle_context.
        - Second reply_stream call with UserConfirmResultEvent: the same
          reply_id resumes. on_reasoning reads 300 >= 300 from middle_context
          and injects the hint + forces tool_choice=none before the final
          model call, proving budget state survived the HITL round-trip.
        """
        tool_call_id = "tc_hitl"
        tool_input = "{}"
        toolkit = Toolkit(tools=[ConfirmRequiredTool()])

        model = MockModel()
        model.set_responses(
            [
                # Step 1: model produces a tool call that requires confirmation
                [
                    ChatResponse(
                        content=[
                            ToolCallBlock(
                                id=tool_call_id,
                                name="confirm_required",
                                input=tool_input,
                            ),
                        ],
                        is_last=True,
                        usage=ChatUsage(
                            input_tokens=200,
                            output_tokens=100,
                            time=0.0,
                        ),
                    ),
                ],
                # Step 2 (after confirmation): final wrap-up text
                [
                    ChatResponse(
                        content=[TextBlock(text="wrap up")],
                        is_last=True,
                        usage=ChatUsage(
                            input_tokens=50,
                            output_tokens=20,
                            time=0.0,
                        ),
                    ),
                ],
            ],
        )

        # 200*1 + 100*1 = 300 == token_budget → reasoning after resume
        # injects hint
        middleware = ReplyBudgetControlMiddleware(token_budget=300)
        agent = Agent(
            name="test_agent",
            system_prompt="you are helpful",
            model=model,
            toolkit=toolkit,
            middlewares=[middleware],
        )

        # --- First call: pauses at REQUIRE_USER_CONFIRM ---
        events = []
        async for event in agent.reply_stream(UserMsg("user", "hello")):
            events.append(event)

        event_types = [e.type for e in events]
        self.assertIn("REQUIRE_USER_CONFIRM", event_types)
        self.assertNotIn("REPLY_END", event_types)

        reply_id = agent.state.reply_id

        # Token count must be stored in middle_context (survived the pause)
        middleware_key = await middleware.get_middleware_key()
        stored = agent.state.middle_context.get(middleware_key, {})
        self.assertAlmostEqual(stored.get(reply_id, 0), 300.0)

        # --- Second call: resume with user confirmation ---
        user_confirm_event = UserConfirmResultEvent(
            reply_id=reply_id,
            confirm_results=[
                ConfirmResult(
                    confirmed=True,
                    tool_call=ToolCallBlock(
                        id=tool_call_id,
                        name="confirm_required",
                        input=tool_input,
                    ),
                ),
            ],
        )

        resume_events = []
        async for event in agent.reply_stream(inputs=user_confirm_event):
            resume_events.append(event)

        resume_event_types = [e.type for e in resume_events]
        self.assertIn("REPLY_END", resume_event_types)

        # Hint is appended to the existing assistant message (which was created
        # in the first call), so we search the full context rather than a
        # slice.
        hint_msgs = [
            m
            for m in agent.state.context
            if _has_hint_block(m, middleware.hint_message)
        ]
        self.assertGreater(len(hint_msgs), 0)

        # middle_context must be cleaned up after reply ends
        bucket = agent.state.middle_context.get(middleware_key, {})
        self.assertNotIn(reply_id, bucket)