# -*- 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)