# -*- coding: utf-8 -*- # pylint: disable=redefined-builtin """Test the user confirmation events in the agent class.""" from typing import Any from unittest.async_case import IsolatedAsyncioTestCase from utils import AnyString, MockModel from agentscope.agent import Agent from agentscope.model import ChatResponse from agentscope.tool import ( ToolBase, Toolkit, ToolChunk, ) from agentscope.permission import ( PermissionDecision, PermissionBehavior, PermissionContext, ) from agentscope.message import ( TextBlock, ToolCallBlock, UserMsg, ) from agentscope.event import UserConfirmResultEvent, ConfirmResult class MockUserConfirmSequentialTool(ToolBase): """A mock tool that requires user confirmation (sequential).""" name: str = "mock_user_confirm_sequential_tool" description: str = "A mock user confirm sequential tool for testing" input_schema: dict[str, Any] = { "type": "object", "properties": { "input": {"type": "string", "description": "Input string"}, }, "required": ["input"], } 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: """Check permissions for the tool usage.""" return PermissionDecision( behavior=PermissionBehavior.ASK, decision_reason="Mock tool requires user confirmation", message="Mock tool requires user confirmation", ) async def __call__(self, input: str, **kwargs: Any) -> ToolChunk: """Execute the tool.""" return ToolChunk( content=[ TextBlock(text=f"User confirm sequential result: {input}"), ], ) class MockUserConfirmConcurrentTool(ToolBase): """A mock tool that requires user confirmation (concurrent).""" name: str = "mock_user_confirm_concurrent_tool" description: str = "A mock user confirm concurrent tool for testing" input_schema: dict[str, Any] = { "type": "object", "properties": { "input": {"type": "string", "description": "Input string"}, }, "required": ["input"], } is_concurrency_safe: bool = True 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: """Check permissions for the tool usage.""" return PermissionDecision( behavior=PermissionBehavior.ASK, decision_reason="Mock tool requires user confirmation", message="Mock tool requires user confirmation", ) async def __call__(self, input: str, **kwargs: Any) -> ToolChunk: """Execute the tool.""" return ToolChunk( content=[ TextBlock(text=f"User confirm concurrent result: {input}"), ], ) class AgentUserConfirmationTest(IsolatedAsyncioTestCase): """Test the user confirmation events in the agent class.""" def _get_tool_call_events( self, id: str, name: str, delta: str, ) -> list[dict]: """Helper method to get the expected tool call events.""" return [ { "type": "TOOL_CALL_START", "tool_call_id": id, "tool_call_name": name, }, { "type": "TOOL_CALL_DELTA", "tool_call_id": id, "delta": delta, }, { "type": "TOOL_CALL_END", "tool_call_id": id, }, ] def _get_tool_result_events( self, id: str, name: str, result: str, ) -> list[dict]: """Helper method to get the expected tool result events.""" return [ { "type": "TOOL_RESULT_START", "tool_call_id": id, "tool_call_name": name, }, { "type": "TOOL_RESULT_TEXT_DELTA", "tool_call_id": id, "delta": result, }, { "type": "TOOL_RESULT_END", "tool_call_id": id, "state": "success", }, ] async def asyncSetUp(self) -> None: """The async setup method.""" self.model = MockModel() self.agent = Agent( name="Friday", system_prompt="You are a helpful assistant.", model=self.model, toolkit=Toolkit(), ) self.tool_call_id_1 = "tool_call_1" self.tool_call_id_2 = "tool_call_2" self.user_input_text = "Test" self.tool_input_1 = '{"input": "test1"}' self.tool_input_2 = '{"input": "test2"}' self.sequential_tool_name = "mock_user_confirm_sequential_tool" self.concurrent_tool_name = "mock_user_confirm_concurrent_tool" self.sequential_result_1 = "User confirm sequential result: test1" self.sequential_result_2 = "User confirm sequential result: test2" self.concurrent_result_1 = "User confirm concurrent result: test1" self.concurrent_result_2 = "User confirm concurrent result: test2" self.final_response_text = "Result 1" self.final_text_events = [ { "type": "MODEL_CALL_START", "model_name": "mock-model", }, { "type": "TEXT_BLOCK_START", "block_id": AnyString(), }, { "type": "TEXT_BLOCK_DELTA", "block_id": AnyString(), "delta": self.final_response_text, }, { "type": "TEXT_BLOCK_END", "block_id": AnyString(), }, { "type": "MODEL_CALL_END", "input_tokens": 0, "output_tokens": 0, }, ] self.final_mock_responses = [ ChatResponse( content=[ TextBlock(text=self.final_response_text), ], is_last=False, ), ChatResponse( content=[ TextBlock(text=self.final_response_text), ], is_last=True, ), ] def _get_event_base(self, reply_id: str) -> dict: """Get the dict with the basic fields for event assertion.""" return { "id": AnyString(), "created_at": AnyString(), "metadata": {}, "reply_id": reply_id, } def _get_msg_base(self) -> dict: """Get the dict with the basic fields for message assertion.""" return { "id": AnyString(), "created_at": AnyString(), "finished_at": None, "metadata": {}, "name": "Friday", "role": "assistant", "usage": None, } async def test_single_user_confirmation(self) -> None: """Test single user confirmation tool call. The agent should: 1. Generate a tool call that requires user confirmation 2. Emit REQUIRE_USER_CONFIRM event and pause 3. Resume when UserConfirmResultEvent is provided 4. Execute the tool and continue """ # Register user confirm tool confirm_tool = MockUserConfirmSequentialTool() self.agent.toolkit = Toolkit( tools=[confirm_tool], ) # Set up mock response with tool call (no final text response) self.model.set_responses( [ [ ChatResponse( content=[ ToolCallBlock( id=self.tool_call_id_1, name=self.sequential_tool_name, input=self.tool_input_1, ), ], is_last=False, ), ChatResponse( content=[ ToolCallBlock( id=self.tool_call_id_1, name=self.sequential_tool_name, input=self.tool_input_1, ), ], is_last=True, ), ], self.final_mock_responses, ], ) # First call: collect events until REQUIRE_USER_CONFIRM events = [] async for event in self.agent.reply_stream( UserMsg(name="user", content=self.user_input_text), ): events.append(event.model_dump()) # Verify events session_id = self.agent.state.session_id reply_id = self.agent.state.reply_id expected_events = [ { "type": "REPLY_START", "session_id": session_id, "name": "Friday", "role": "assistant", }, {"type": "MODEL_CALL_START", "model_name": "mock-model"}, *self._get_tool_call_events( self.tool_call_id_1, self.sequential_tool_name, self.tool_input_1, ), { "type": "MODEL_CALL_END", "input_tokens": 0, "output_tokens": 0, }, { "type": "REQUIRE_USER_CONFIRM", "reply_id": reply_id, "tool_calls": [ { "type": "tool_call", "id": self.tool_call_id_1, "name": self.sequential_tool_name, "input": self.tool_input_1, "state": "asking", "suggested_rules": [ { "tool_name": self.sequential_tool_name, "rule_content": None, "behavior": PermissionBehavior.ALLOW, "source": "suggested", }, ], }, ], }, ] basic_dict = self._get_event_base(reply_id) self.assertListEqual( events, [{**basic_dict, **_} for _ in expected_events], ) # Assert context after first call msg_base = self._get_msg_base() expected_context = [ { "name": "user", "role": "user", "content": [ { "type": "text", "id": AnyString(), "text": self.user_input_text, }, ], "finished_at": AnyString(), }, { "content": [ { "type": "tool_call", "id": self.tool_call_id_1, "name": self.sequential_tool_name, "input": self.tool_input_1, "state": "asking", "suggested_rules": [ { "tool_name": self.sequential_tool_name, "rule_content": None, "behavior": PermissionBehavior.ALLOW, "source": "suggested", }, ], }, ], }, ] context_dicts = [msg.model_dump() for msg in self.agent.state.context] expected_context = [{**msg_base, **_} for _ in expected_context] self.assertListEqual(context_dicts, expected_context) # Create user confirmation result event user_confirm_event = UserConfirmResultEvent( reply_id=reply_id, confirm_results=[ ConfirmResult( confirmed=True, tool_call=ToolCallBlock( id=self.tool_call_id_1, name=self.sequential_tool_name, input=self.tool_input_1, ), ), ], ) # Second call: resume with user confirmation result events = [] async for event in self.agent.reply_stream(inputs=user_confirm_event): events.append(event.model_dump()) # Verify events after resumption expected_events_resume = [ *self._get_tool_result_events( self.tool_call_id_1, self.sequential_tool_name, self.sequential_result_1, ), *self.final_text_events, {"type": "REPLY_END", "session_id": session_id}, ] self.assertListEqual( events, [{**basic_dict, **_} for _ in expected_events_resume], ) # Assert final context expected_context_final = [ { "name": "user", "role": "user", "content": [ { "type": "text", "id": AnyString(), "text": self.user_input_text, }, ], "finished_at": AnyString(), }, { "content": [ { "type": "tool_call", "id": self.tool_call_id_1, "name": self.sequential_tool_name, "input": self.tool_input_1, "state": "finished", "suggested_rules": [ { "tool_name": self.sequential_tool_name, "rule_content": None, "behavior": PermissionBehavior.ALLOW, "source": "suggested", }, ], }, { "type": "tool_result", "id": AnyString(), "name": self.sequential_tool_name, "output": [ { "type": "text", "id": AnyString(), "text": self.sequential_result_1, }, ], "state": "success", "metadata": {}, }, { "type": "text", "id": AnyString(), "text": self.final_response_text, }, ], }, ] context_dicts = [msg.model_dump() for msg in self.agent.state.context] expected_context_final = [ {**msg_base, **_} for _ in expected_context_final ] self.assertListEqual(context_dicts, expected_context_final) async def test_sequential_user_confirmation(self) -> None: """Test multiple user confirmation tool calls in sequential execution. The agent should: 1. Generate multiple tool calls that require user confirmation 2. All tools have is_concurrent_safe=False (sequential) 3. Emit REQUIRE_USER_CONFIRM event and pause 4. Resume when UserConfirmResultEvent is provided 5. Execute the tools and continue """ # Register user confirm sequential tool confirm_tool = MockUserConfirmSequentialTool() self.agent.toolkit = Toolkit( tools=[confirm_tool], ) # Set up mock response with multiple tool calls self.model.set_responses( [ [ ChatResponse( content=[ ToolCallBlock( id=self.tool_call_id_1, name=self.sequential_tool_name, input=self.tool_input_1, ), ToolCallBlock( id=self.tool_call_id_2, name=self.sequential_tool_name, input=self.tool_input_2, ), ], is_last=False, usage=None, ), ChatResponse( content=[ ToolCallBlock( id=self.tool_call_id_1, name=self.sequential_tool_name, input=self.tool_input_1, ), ToolCallBlock( id=self.tool_call_id_2, name=self.sequential_tool_name, input=self.tool_input_2, ), ], is_last=True, usage=None, ), ], self.final_mock_responses, ], ) # First call: collect events until REQUIRE_USER_CONFIRM events = [] async for event in self.agent.reply_stream( UserMsg(name="user", content=self.user_input_text), ): events.append(event.model_dump()) # Verify events session_id = self.agent.state.session_id reply_id = self.agent.state.reply_id tool_call_1_events = self._get_tool_call_events( self.tool_call_id_1, self.sequential_tool_name, self.tool_input_1, ) tool_call_2_events = self._get_tool_call_events( self.tool_call_id_2, self.sequential_tool_name, self.tool_input_2, ) expected_events = [ { "type": "REPLY_START", "session_id": session_id, "name": "Friday", "role": "assistant", }, {"type": "MODEL_CALL_START", "model_name": "mock-model"}, *tool_call_1_events[:2], *tool_call_2_events[:2], tool_call_1_events[2], tool_call_2_events[2], { "type": "MODEL_CALL_END", "input_tokens": 0, "output_tokens": 0, }, { "type": "REQUIRE_USER_CONFIRM", "reply_id": reply_id, "tool_calls": [ { "type": "tool_call", "id": self.tool_call_id_1, "name": self.sequential_tool_name, "input": self.tool_input_1, "state": "asking", "suggested_rules": [ { "tool_name": self.sequential_tool_name, "rule_content": None, "behavior": PermissionBehavior.ALLOW, "source": "suggested", }, ], }, ], }, ] basic_dict = self._get_event_base(reply_id) self.assertListEqual( events, [{**basic_dict, **_} for _ in expected_events], ) # Assert context after first call msg_base = self._get_msg_base() expected_context = [ { "name": "user", "role": "user", "content": [ { "type": "text", "id": AnyString(), "text": self.user_input_text, }, ], "finished_at": AnyString(), }, { "content": [ { "type": "tool_call", "id": self.tool_call_id_1, "name": self.sequential_tool_name, "input": self.tool_input_1, "state": "asking", "suggested_rules": [ { "tool_name": self.sequential_tool_name, "rule_content": None, "behavior": PermissionBehavior.ALLOW, "source": "suggested", }, ], }, { "type": "tool_call", "id": self.tool_call_id_2, "name": self.sequential_tool_name, "input": self.tool_input_2, "state": "pending", "suggested_rules": [], }, ], }, ] context_dicts = [msg.model_dump() for msg in self.agent.state.context] expected_context = [{**msg_base, **_} for _ in expected_context] self.assertListEqual(context_dicts, expected_context) # Create user confirmation result event user_confirm_event = UserConfirmResultEvent( reply_id=reply_id, confirm_results=[ ConfirmResult( confirmed=True, tool_call=ToolCallBlock( id=self.tool_call_id_1, name=self.sequential_tool_name, input=self.tool_input_1, ), ), ], ) # resume with user confirmation result events = [] async for event in self.agent.reply_stream(inputs=user_confirm_event): events.append(event.model_dump()) # Verify events after resumption (sequential execution) expected_events_resume = [ *self._get_tool_result_events( self.tool_call_id_1, self.sequential_tool_name, self.sequential_result_1, ), { "type": "REQUIRE_USER_CONFIRM", "tool_calls": [ { "type": "tool_call", "id": self.tool_call_id_2, "name": self.sequential_tool_name, "input": self.tool_input_2, "state": "asking", "suggested_rules": [ { "tool_name": self.sequential_tool_name, "rule_content": None, "behavior": PermissionBehavior.ALLOW, "source": "suggested", }, ], }, ], }, ] self.assertListEqual( events, [{**basic_dict, **_} for _ in expected_events_resume], ) # Confirm the second tool call user_confirm_event = UserConfirmResultEvent( reply_id=reply_id, confirm_results=[ ConfirmResult( confirmed=True, tool_call=ToolCallBlock( id=self.tool_call_id_2, name=self.sequential_tool_name, input=self.tool_input_2, ), ), ], ) # Second call: resume with user confirmation result events = [] async for event in self.agent.reply_stream(inputs=user_confirm_event): events.append(event.model_dump()) expected_events_resume_2 = [ *self._get_tool_result_events( self.tool_call_id_2, self.sequential_tool_name, self.sequential_result_2, ), *self.final_text_events, {"type": "REPLY_END", "session_id": session_id}, ] self.assertListEqual( events, [{**basic_dict, **_} for _ in expected_events_resume_2], ) # Assert final context expected_context_final = [ { "name": "user", "role": "user", "content": [ { "type": "text", "id": AnyString(), "text": self.user_input_text, }, ], "finished_at": AnyString(), }, { "content": [ { "type": "tool_call", "id": self.tool_call_id_1, "name": self.sequential_tool_name, "input": self.tool_input_1, "state": "finished", "suggested_rules": [ { "tool_name": self.sequential_tool_name, "rule_content": None, "behavior": PermissionBehavior.ALLOW, "source": "suggested", }, ], }, { "type": "tool_call", "id": self.tool_call_id_2, "name": self.sequential_tool_name, "input": self.tool_input_2, "state": "finished", "suggested_rules": [ { "tool_name": self.sequential_tool_name, "rule_content": None, "behavior": PermissionBehavior.ALLOW, "source": "suggested", }, ], }, { "type": "tool_result", "id": AnyString(), "name": self.sequential_tool_name, "output": [ { "type": "text", "id": AnyString(), "text": self.sequential_result_1, }, ], "state": "success", "metadata": {}, }, { "type": "tool_result", "id": AnyString(), "name": self.sequential_tool_name, "output": [ { "type": "text", "id": AnyString(), "text": self.sequential_result_2, }, ], "state": "success", "metadata": {}, }, { "type": "text", "id": AnyString(), "text": self.final_response_text, }, ], }, ] context_dicts = [msg.model_dump() for msg in self.agent.state.context] expected_context_final = [ {**msg_base, **_} for _ in expected_context_final ] self.assertListEqual(context_dicts, expected_context_final) async def test_concurrent_user_confirmation(self) -> None: """Test multiple user confirmation tool calls in concurrent execution. The agent should: 1. Generate multiple tool calls that require user confirmation 2. All tools have is_concurrent_safe=True (concurrent) 3. Emit REQUIRE_USER_CONFIRM event and pause 4. Resume when UserConfirmResultEvent is provided 5. Execute the tools and continue """ # Register user confirm concurrent tool confirm_tool = MockUserConfirmConcurrentTool() self.agent.toolkit = Toolkit( tools=[confirm_tool], ) # Set up mock response with multiple tool calls self.model.set_responses( [ [ ChatResponse( content=[ ToolCallBlock( id=self.tool_call_id_1, name=self.concurrent_tool_name, input=self.tool_input_1, ), ToolCallBlock( id=self.tool_call_id_2, name=self.concurrent_tool_name, input=self.tool_input_2, ), ], is_last=False, ), ChatResponse( content=[ ToolCallBlock( id=self.tool_call_id_1, name=self.concurrent_tool_name, input=self.tool_input_1, ), ToolCallBlock( id=self.tool_call_id_2, name=self.concurrent_tool_name, input=self.tool_input_2, ), ], is_last=True, ), ], self.final_mock_responses, ], ) # First call: collect events until REQUIRE_USER_CONFIRM events = [] async for event in self.agent.reply_stream( UserMsg(name="user", content=self.user_input_text), ): events.append(event.model_dump()) # Verify events session_id = self.agent.state.session_id reply_id = self.agent.state.reply_id tool_call_1_events = self._get_tool_call_events( self.tool_call_id_1, self.concurrent_tool_name, self.tool_input_1, ) tool_call_2_events = self._get_tool_call_events( self.tool_call_id_2, self.concurrent_tool_name, self.tool_input_2, ) expected_events = [ { "type": "REPLY_START", "session_id": session_id, "name": "Friday", "role": "assistant", }, {"type": "MODEL_CALL_START", "model_name": "mock-model"}, *tool_call_1_events[:2], *tool_call_2_events[:2], tool_call_1_events[2], tool_call_2_events[2], { "type": "MODEL_CALL_END", "input_tokens": 0, "output_tokens": 0, }, { "type": "REQUIRE_USER_CONFIRM", "reply_id": reply_id, "tool_calls": [ { "type": "tool_call", "id": self.tool_call_id_1, "name": self.concurrent_tool_name, "input": self.tool_input_1, "state": "asking", "suggested_rules": [ { "tool_name": self.concurrent_tool_name, "rule_content": None, "behavior": PermissionBehavior.ALLOW, "source": "suggested", }, ], }, ], }, { "type": "REQUIRE_USER_CONFIRM", "reply_id": reply_id, "tool_calls": [ { "type": "tool_call", "id": self.tool_call_id_2, "name": self.concurrent_tool_name, "input": self.tool_input_2, "state": "asking", "suggested_rules": [ { "tool_name": self.concurrent_tool_name, "rule_content": None, "behavior": PermissionBehavior.ALLOW, "source": "suggested", }, ], }, ], }, ] basic_dict = self._get_event_base(reply_id) self.assertListEqual( events, [{**basic_dict, **_} for _ in expected_events], ) # Assert context after first call msg_base = self._get_msg_base() expected_context = [ { "name": "user", "role": "user", "content": [ { "type": "text", "id": AnyString(), "text": self.user_input_text, }, ], "finished_at": AnyString(), }, { "content": [ { "type": "tool_call", "id": self.tool_call_id_1, "name": self.concurrent_tool_name, "input": self.tool_input_1, "state": "asking", "suggested_rules": [ { "tool_name": self.concurrent_tool_name, "rule_content": None, "behavior": PermissionBehavior.ALLOW, "source": "suggested", }, ], }, { "type": "tool_call", "id": self.tool_call_id_2, "name": self.concurrent_tool_name, "input": self.tool_input_2, "state": "asking", "suggested_rules": [ { "tool_name": self.concurrent_tool_name, "rule_content": None, "behavior": PermissionBehavior.ALLOW, "source": "suggested", }, ], }, ], }, ] context_dicts = [msg.model_dump() for msg in self.agent.state.context] expected_context = [{**msg_base, **_} for _ in expected_context] self.assertListEqual(context_dicts, expected_context) # Create user confirmation result event user_confirm_event = UserConfirmResultEvent( reply_id=reply_id, confirm_results=[ ConfirmResult( confirmed=True, tool_call=ToolCallBlock( id=self.tool_call_id_1, name=self.concurrent_tool_name, input=self.tool_input_1, ), ), ], ) # resume with user confirmation result events = [] async for event in self.agent.reply_stream(inputs=user_confirm_event): events.append(event.model_dump()) # Verify events for tool call 1 after resumption expected_events = [ *self._get_tool_result_events( self.tool_call_id_1, self.concurrent_tool_name, self.concurrent_result_1, ), ] self.assertListEqual( events, [{**basic_dict, **_} for _ in expected_events], ) # The second tool call user_confirm_event = UserConfirmResultEvent( reply_id=reply_id, confirm_results=[ ConfirmResult( confirmed=True, tool_call=ToolCallBlock( id=self.tool_call_id_2, name=self.concurrent_tool_name, input=self.tool_input_2, ), ), ], ) events = [] async for event in self.agent.reply_stream(inputs=user_confirm_event): events.append(event.model_dump()) expected_events = [ *self._get_tool_result_events( self.tool_call_id_2, self.concurrent_tool_name, self.concurrent_result_2, ), *self.final_text_events, {"type": "REPLY_END", "session_id": session_id}, ] self.assertListEqual( events, [{**basic_dict, **_} for _ in expected_events], ) # Assert final context expected_context_final = [ { "name": "user", "role": "user", "content": [ { "type": "text", "id": AnyString(), "text": self.user_input_text, }, ], "finished_at": AnyString(), }, { "content": [ { "type": "tool_call", "id": self.tool_call_id_1, "name": self.concurrent_tool_name, "input": self.tool_input_1, "state": "finished", "suggested_rules": [ { "tool_name": self.concurrent_tool_name, "rule_content": None, "behavior": PermissionBehavior.ALLOW, "source": "suggested", }, ], }, { "type": "tool_call", "id": self.tool_call_id_2, "name": self.concurrent_tool_name, "input": self.tool_input_2, "state": "finished", "suggested_rules": [ { "tool_name": self.concurrent_tool_name, "rule_content": None, "behavior": PermissionBehavior.ALLOW, "source": "suggested", }, ], }, { "type": "tool_result", "id": AnyString(), "name": self.concurrent_tool_name, "output": [ { "type": "text", "id": AnyString(), "text": self.concurrent_result_1, }, ], "state": "success", "metadata": {}, }, { "type": "tool_result", "id": AnyString(), "name": self.concurrent_tool_name, "output": [ { "type": "text", "id": AnyString(), "text": self.concurrent_result_2, }, ], "state": "success", "metadata": {}, }, { "type": "text", "id": AnyString(), "text": self.final_response_text, }, ], }, ] context_dicts = [msg.model_dump() for msg in self.agent.state.context] expected_context_final = [ {**msg_base, **_} for _ in expected_context_final ] self.assertListEqual(context_dicts, expected_context_final) async def test_concurrent_user_confirmation_in_single_event(self) -> None: """Test concurrent user confirmation when two approvals arrive together. The agent should: 1. Generate multiple tool calls that require user confirmation 2. Pause in concurrent mode with two asking tool calls 3. Resume when one UserConfirmResultEvent carries both confirmations 4. Execute both tools and continue reasoning after both complete """ confirm_tool = MockUserConfirmConcurrentTool() self.agent.toolkit = Toolkit( tools=[confirm_tool], ) self.model.set_responses( [ [ ChatResponse( content=[ ToolCallBlock( id=self.tool_call_id_1, name=self.concurrent_tool_name, input=self.tool_input_1, ), ToolCallBlock( id=self.tool_call_id_2, name=self.concurrent_tool_name, input=self.tool_input_2, ), ], is_last=False, ), ChatResponse( content=[ ToolCallBlock( id=self.tool_call_id_1, name=self.concurrent_tool_name, input=self.tool_input_1, ), ToolCallBlock( id=self.tool_call_id_2, name=self.concurrent_tool_name, input=self.tool_input_2, ), ], is_last=True, ), ], self.final_mock_responses, ], ) events = [] async for event in self.agent.reply_stream( UserMsg(name="user", content=self.user_input_text), ): events.append(event.model_dump()) session_id = self.agent.state.session_id reply_id = self.agent.state.reply_id basic_dict = self._get_event_base(reply_id) msg_base = self._get_msg_base() tool_call_1_events = self._get_tool_call_events( self.tool_call_id_1, self.concurrent_tool_name, self.tool_input_1, ) tool_call_2_events = self._get_tool_call_events( self.tool_call_id_2, self.concurrent_tool_name, self.tool_input_2, ) expected_events = [ { "type": "REPLY_START", "session_id": session_id, "name": "Friday", "role": "assistant", }, {"type": "MODEL_CALL_START", "model_name": "mock-model"}, *tool_call_1_events[:2], *tool_call_2_events[:2], tool_call_1_events[2], tool_call_2_events[2], { "type": "MODEL_CALL_END", "input_tokens": 0, "output_tokens": 0, }, { "type": "REQUIRE_USER_CONFIRM", "reply_id": reply_id, "tool_calls": [ { "type": "tool_call", "id": self.tool_call_id_1, "name": self.concurrent_tool_name, "input": self.tool_input_1, "state": "asking", "suggested_rules": [ { "tool_name": self.concurrent_tool_name, "rule_content": None, "behavior": PermissionBehavior.ALLOW, "source": "suggested", }, ], }, ], }, { "type": "REQUIRE_USER_CONFIRM", "reply_id": reply_id, "tool_calls": [ { "type": "tool_call", "id": self.tool_call_id_2, "name": self.concurrent_tool_name, "input": self.tool_input_2, "state": "asking", "suggested_rules": [ { "tool_name": self.concurrent_tool_name, "rule_content": None, "behavior": PermissionBehavior.ALLOW, "source": "suggested", }, ], }, ], }, ] self.assertListEqual( events, [{**basic_dict, **_} for _ in expected_events], ) expected_context = [ { "name": "user", "role": "user", "content": [ { "type": "text", "id": AnyString(), "text": self.user_input_text, }, ], "finished_at": AnyString(), }, { "content": [ { "type": "tool_call", "id": self.tool_call_id_1, "name": self.concurrent_tool_name, "input": self.tool_input_1, "state": "asking", "suggested_rules": [ { "tool_name": self.concurrent_tool_name, "rule_content": None, "behavior": PermissionBehavior.ALLOW, "source": "suggested", }, ], }, { "type": "tool_call", "id": self.tool_call_id_2, "name": self.concurrent_tool_name, "input": self.tool_input_2, "state": "asking", "suggested_rules": [ { "tool_name": self.concurrent_tool_name, "rule_content": None, "behavior": PermissionBehavior.ALLOW, "source": "suggested", }, ], }, ], }, ] context_dicts = [msg.model_dump() for msg in self.agent.state.context] expected_context = [{**msg_base, **_} for _ in expected_context] self.assertListEqual(context_dicts, expected_context) user_confirm_event = UserConfirmResultEvent( reply_id=reply_id, confirm_results=[ ConfirmResult( confirmed=True, tool_call=ToolCallBlock( id=self.tool_call_id_1, name=self.concurrent_tool_name, input=self.tool_input_1, ), ), ConfirmResult( confirmed=True, tool_call=ToolCallBlock( id=self.tool_call_id_2, name=self.concurrent_tool_name, input=self.tool_input_2, ), ), ], ) events = [] async for event in self.agent.reply_stream(inputs=user_confirm_event): events.append(event.model_dump()) tool_events = events[:6] final_events = events[6:] self.assertEqual(len(tool_events), 6) expected_tool_events = { self.tool_call_id_1: [ {**basic_dict, **_} for _ in self._get_tool_result_events( self.tool_call_id_1, self.concurrent_tool_name, self.concurrent_result_1, ) ], self.tool_call_id_2: [ {**basic_dict, **_} for _ in self._get_tool_result_events( self.tool_call_id_2, self.concurrent_tool_name, self.concurrent_result_2, ) ], } for tool_call_id, expected_tool_event in expected_tool_events.items(): self.assertListEqual( [ event for event in tool_events if event["tool_call_id"] == tool_call_id ], expected_tool_event, ) expected_final_events = [ *self.final_text_events, {"type": "REPLY_END", "session_id": session_id}, ] self.assertListEqual( final_events, [{**basic_dict, **_} for _ in expected_final_events], ) self.assertEqual(len(self.agent.state.context), 2) self.assertEqual( self.agent.state.context[0].model_dump(), { "id": AnyString(), "created_at": AnyString(), "finished_at": AnyString(), "metadata": {}, "name": "user", "role": "user", "content": [ { "type": "text", "id": AnyString(), "text": self.user_input_text, }, ], "usage": None, }, ) assistant_msg = self.agent.state.context[-1] self.assertEqual( [ _.model_dump()["state"] for _ in assistant_msg.get_content_blocks("tool_call") ], ["finished", "finished"], ) self.assertEqual( [ _.model_dump()["id"] for _ in assistant_msg.get_content_blocks("tool_call") ], [self.tool_call_id_1, self.tool_call_id_2], ) self.assertEqual( { ( _.model_dump()["name"], _.model_dump()["state"], _.output[0].text, ) for _ in assistant_msg.get_content_blocks("tool_result") }, { ( self.concurrent_tool_name, "success", self.concurrent_result_1, ), ( self.concurrent_tool_name, "success", self.concurrent_result_2, ), }, ) self.assertEqual( [_.text for _ in assistant_msg.get_content_blocks("text")], [self.final_response_text], ) async def asyncTearDown(self) -> None: """The async teardown method."""