# -*- coding: utf-8 -*- """The basic test of 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 class MockSequentialTool(ToolBase): """A mock tool that is not concurrency safe (sequential execution).""" name: str = "mock_sequential_tool" description: str = "A mock 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 = True 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.ALLOW, decision_reason="Mock tool always allows", message="Mock tool always allows", ) # pylint: disable=redefined-builtin async def __call__(self, input: str, **kwargs: Any) -> ToolChunk: """Execute the tool.""" return ToolChunk( content=[TextBlock(text=f"Sequential result: {input}")], ) class MockConcurrentTool(ToolBase): """A mock tool that is concurrency safe (concurrent execution).""" name: str = "mock_concurrent_tool" description: str = "A mock 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 = True 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.ALLOW, decision_reason="Mock tool always allows", message="Mock tool always allows", ) # pylint: disable=redefined-builtin async def __call__(self, input: str, **kwargs: Any) -> ToolChunk: """Execute the tool.""" return ToolChunk( content=[TextBlock(text=f"Concurrent result: {input}")], ) class AgentBasicTest(IsolatedAsyncioTestCase): """The basic test of the agent class.""" 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(), ) 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_default_configs_are_not_shared_between_agents( self, ) -> None: """Agents created with defaults own independent config objects.""" agent_1 = Agent( name="agent-1", system_prompt="You are agent 1.", model=MockModel(), ) agent_2 = Agent( name="agent-2", system_prompt="You are agent 2.", model=MockModel(), ) self.assertIsNot(agent_1.model_config, agent_2.model_config) self.assertIsNot(agent_1.context_config, agent_2.context_config) self.assertIsNot(agent_1.react_config, agent_2.react_config) agent_1.model_config.max_retries = 3 agent_1.context_config.tool_result_limit = 123 agent_1.react_config.max_iters = 2 self.assertNotEqual( agent_1.model_config.max_retries, agent_2.model_config.max_retries, ) self.assertNotEqual( agent_1.context_config.tool_result_limit, agent_2.context_config.tool_result_limit, ) self.assertNotEqual( agent_1.react_config.max_iters, agent_2.react_config.max_iters, ) async def test_streaming_reasoning(self) -> None: """Test the streaming model inference without tool calls generated, only text in model response. Test both the reply and replyStream interfaces. """ # Set up mock responses for streaming self.model.set_responses( [ [ ChatResponse( content=[TextBlock(text="Hello")], is_last=False, ), ChatResponse( content=[TextBlock(text=" world")], is_last=False, ), ChatResponse(content=[TextBlock(text="!")], is_last=False), ChatResponse( content=[TextBlock(text="Hello world!")], is_last=True, ), ], ], ) # Test replyStream interface events = [] async for event in self.agent.reply_stream( UserMsg(name="user", content="Hi"), ): 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", }, { "type": "TEXT_BLOCK_START", "block_id": AnyString(), }, { "type": "TEXT_BLOCK_DELTA", "block_id": AnyString(), "delta": "Hello", }, { "type": "TEXT_BLOCK_DELTA", "block_id": AnyString(), "delta": " world", }, { "type": "TEXT_BLOCK_DELTA", "block_id": AnyString(), "delta": "!", }, { "type": "TEXT_BLOCK_END", "block_id": AnyString(), }, { "type": "MODEL_CALL_END", "input_tokens": 0, "output_tokens": 0, "finished_reason": "completed", }, { "type": "REPLY_END", "session_id": session_id, "finished_reason": "completed", }, ] basic_dict = self._get_event_base(reply_id) self.assertListEqual( events, [{**basic_dict, **_} for _ in expected_events], ) # Assert context after reply_stream msg_base = self._get_msg_base() expected_context = [ { **msg_base, "name": "user", "role": "user", "content": [ { "type": "text", "id": AnyString(), "text": "Hi", }, ], "finished_at": AnyString(), }, { **msg_base, "content": [ { "type": "text", "id": AnyString(), "text": "Hello world!", }, ], }, ] context_dicts = [msg.model_dump() for msg in self.agent.state.context] self.assertListEqual(context_dicts, expected_context) # Test reply interface self.model.cnt = 0 # Reset mock model response index msg = await self.agent.reply(UserMsg(name="user", content="Hi again")) self.assertDictEqual( msg.model_dump(), { "name": "Friday", "role": "assistant", "content": [ { "type": "text", "id": AnyString(), "text": "Hello world!", }, ], "created_at": AnyString(), "finished_at": None, "id": AnyString(), "metadata": {}, "usage": None, }, ) # Assert context after reply expected_context_after_reply = [ { **msg_base, "name": "user", "role": "user", "content": [ { "type": "text", "id": AnyString(), "text": "Hi", }, ], "metadata": {}, "finished_at": AnyString(), }, { **msg_base, "name": "Friday", "role": "assistant", "content": [ { "type": "text", "id": AnyString(), "text": "Hello world!", }, ], "metadata": {}, }, { **msg_base, "name": "user", "role": "user", "content": [ { "type": "text", "id": AnyString(), "text": "Hi again", }, ], "metadata": {}, "finished_at": AnyString(), }, { **msg_base, "name": "Friday", "role": "assistant", "content": [ { "type": "text", "id": AnyString(), "text": "Hello world!", }, ], "metadata": {}, }, ] context_dicts = [msg.model_dump() for msg in self.agent.state.context] self.assertListEqual(context_dicts, expected_context_after_reply) async def test_non_streaming_reasoning(self) -> None: """Test the non-streaming model inference without tool calls generated, only text in model response. Test both the reply and replyStream interfaces. """ # Set up mock response for non-streaming self.model.set_responses( [ ChatResponse( content=[TextBlock(text="Hello world!")], is_last=True, usage=None, ), ], ) # Test replyStream interface events = [] async for event in self.agent.reply_stream( UserMsg(name="user", content="Hi"), ): 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", }, { "type": "TEXT_BLOCK_START", "block_id": AnyString(), }, { "type": "TEXT_BLOCK_DELTA", "block_id": AnyString(), "delta": "Hello world!", }, { "type": "TEXT_BLOCK_END", "block_id": AnyString(), }, { "type": "MODEL_CALL_END", "input_tokens": 0, "output_tokens": 0, "finished_reason": "completed", }, { "type": "REPLY_END", "session_id": session_id, "finished_reason": "completed", }, ] basic_dict = self._get_event_base(reply_id) self.assertListEqual( events, [{**basic_dict, **_} for _ in expected_events], ) # Assert context after reply_stream msg_base = self._get_msg_base() expected_context = [ { **msg_base, "name": "user", "role": "user", "content": [ { "type": "text", "id": AnyString(), "text": "Hi", }, ], "metadata": {}, "finished_at": AnyString(), }, { **msg_base, "content": [ { "type": "text", "id": AnyString(), "text": "Hello world!", }, ], "metadata": {}, }, ] context_dicts = [msg.model_dump() for msg in self.agent.state.context] self.assertListEqual(context_dicts, expected_context) # Test reply interface self.model.set_responses( [ ChatResponse( content=[TextBlock(text="Hello world!")], is_last=True, usage=None, ), ], ) msg = await self.agent.reply(UserMsg(name="user", content="Hi again")) self.assertDictEqual( msg.model_dump(), { "name": "Friday", "role": "assistant", "content": [ { "type": "text", "id": AnyString(), "text": "Hello world!", }, ], "created_at": AnyString(), "finished_at": None, "id": AnyString(), "metadata": {}, "usage": None, }, ) # Assert context after reply expected_context_after_reply = [ { **msg_base, "name": "user", "role": "user", "content": [ { "type": "text", "id": AnyString(), "text": "Hi", }, ], "metadata": {}, "finished_at": AnyString(), }, { **msg_base, "name": "Friday", "role": "assistant", "content": [ { "type": "text", "id": AnyString(), "text": "Hello world!", }, ], "metadata": {}, }, { **msg_base, "name": "user", "role": "user", "content": [ { "type": "text", "id": AnyString(), "text": "Hi again", }, ], "metadata": {}, "finished_at": AnyString(), }, { **msg_base, "name": "Friday", "role": "assistant", "content": [ { "type": "text", "id": AnyString(), "text": "Hello world!", }, ], "metadata": {}, }, ] context_dicts = [msg.model_dump() for msg in self.agent.state.context] self.assertListEqual(context_dicts, expected_context_after_reply) async def test_streaming_sequential_tool_calls(self) -> None: """Test the streaming model inference with tool calls generated. Test only the replyStream interface. The `is_concurrent_safe` of the registered tools should be False to make sure the tool calls are executed sequentially. To assert: 1. The events (by assert the dict generated by model_dump) 2. The final reply message 3. The agent state (Before and after replyStream) """ # Register sequential tools seq_tool = MockSequentialTool() self.agent.toolkit = Toolkit(tools=[seq_tool]) # Create tool call IDs tool_call_id_1 = "tool_call_1" tool_call_id_2 = "tool_call_2" # Set up mock responses with tool calls text_block = TextBlock(text="I'll call the tool") tool_call_1 = ToolCallBlock( id=tool_call_id_1, name="mock_sequential_tool", input='{"input": "test1"}', ) tool_call_1_part1 = ToolCallBlock( id=tool_call_id_1, name="mock_sequential_tool", input='{"input": ', ) tool_call_1_part2 = ToolCallBlock( id=tool_call_id_1, name="mock_sequential_tool", input='"test1"}', ) tool_call_2 = ToolCallBlock( id=tool_call_id_2, name="mock_sequential_tool", input='{"input": "test2"}', ) self.model.set_responses( [ [ ChatResponse( content=[text_block, tool_call_1_part1], is_last=False, ), ChatResponse( content=[tool_call_1_part2], is_last=False, ), ChatResponse( content=[tool_call_2], is_last=False, ), ChatResponse( content=[text_block, tool_call_1, tool_call_2], is_last=True, ), ], [ ChatResponse( content=[TextBlock(text="ended")], is_last=False, ), ChatResponse( content=[TextBlock(text="ended")], is_last=True, ), ], ], ) # Collect all events events = [] async for event in self.agent.reply_stream( UserMsg(name="user", content="Test"), ): events.append(event.model_dump(mode="json")) # Verify events session_id = self.agent.state.session_id reply_id = self.agent.state.reply_id # Expected events for sequential tool calls expected_events = [ { "type": "REPLY_START", "session_id": session_id, "name": "Friday", "role": "assistant", }, {"type": "MODEL_CALL_START", "model_name": "mock-model"}, {"type": "TEXT_BLOCK_START", "block_id": AnyString()}, { "type": "TEXT_BLOCK_DELTA", "block_id": AnyString(), "delta": "I'll call the tool", }, { "type": "TOOL_CALL_START", "tool_call_id": tool_call_id_1, "tool_call_name": "mock_sequential_tool", }, { "type": "TOOL_CALL_DELTA", "tool_call_id": tool_call_id_1, "delta": '{"input": ', }, {"type": "TEXT_BLOCK_END", "block_id": AnyString()}, { "type": "TOOL_CALL_DELTA", "tool_call_id": tool_call_id_1, "delta": '"test1"}', }, { "type": "TOOL_CALL_START", "tool_call_id": tool_call_id_2, "tool_call_name": "mock_sequential_tool", }, { "type": "TOOL_CALL_DELTA", "tool_call_id": tool_call_id_2, "delta": '{"input": "test2"}', }, {"type": "TOOL_CALL_END", "tool_call_id": tool_call_id_1}, {"type": "TOOL_CALL_END", "tool_call_id": tool_call_id_2}, { "type": "MODEL_CALL_END", "input_tokens": 0, "output_tokens": 0, "finished_reason": "completed", }, { "type": "TOOL_RESULT_START", "tool_call_id": tool_call_id_1, "tool_call_name": "mock_sequential_tool", }, { "type": "TOOL_RESULT_TEXT_DELTA", "tool_call_id": tool_call_id_1, "delta": "Sequential result: test1", }, { "type": "TOOL_RESULT_END", "tool_call_id": tool_call_id_1, "state": "success", }, { "type": "TOOL_RESULT_START", "tool_call_id": tool_call_id_2, "tool_call_name": "mock_sequential_tool", }, { "type": "TOOL_RESULT_TEXT_DELTA", "tool_call_id": tool_call_id_2, "delta": "Sequential result: test2", }, { "type": "TOOL_RESULT_END", "tool_call_id": tool_call_id_2, "state": "success", }, {"type": "MODEL_CALL_START", "model_name": "mock-model"}, {"type": "TEXT_BLOCK_START", "block_id": AnyString()}, { "type": "TEXT_BLOCK_DELTA", "block_id": AnyString(), "delta": "ended", }, {"type": "TEXT_BLOCK_END", "block_id": AnyString()}, { "type": "MODEL_CALL_END", "input_tokens": 0, "output_tokens": 0, "finished_reason": "completed", }, { "type": "REPLY_END", "session_id": session_id, "finished_reason": "completed", }, ] basic_dict = self._get_event_base(reply_id) self.assertListEqual( events, [{**basic_dict, **_} for _ in expected_events], ) # Assert context after reply_stream msg_base = self._get_msg_base() expected_context = [ { "name": "user", "role": "user", "content": [ { "type": "text", "id": AnyString(), "text": "Test", }, ], "finished_at": AnyString(), }, { "content": [ { "type": "text", "id": AnyString(), "text": "I'll call the tool", }, { "type": "tool_call", "id": tool_call_id_1, "name": "mock_sequential_tool", "input": '{"input": "test1"}', "state": "finished", "suggested_rules": [], }, { "type": "tool_call", "id": tool_call_id_2, "name": "mock_sequential_tool", "input": '{"input": "test2"}', "state": "finished", "suggested_rules": [], }, { "type": "tool_result", "id": AnyString(), "output": [ { "type": "text", "id": AnyString(), "text": "Sequential result: test1", }, ], "name": "mock_sequential_tool", "state": "success", "metadata": {}, }, { "type": "tool_result", "id": AnyString(), "output": [ { "type": "text", "id": AnyString(), "text": "Sequential result: test2", }, ], "name": "mock_sequential_tool", "state": "success", "metadata": {}, }, { "type": "text", "id": AnyString(), "text": "ended", }, ], }, ] 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) async def test_streaming_concurrent_tool_calls(self) -> None: """Test the streaming model inference with tool calls generated. Test only the replyStream interface. The `is_concurrent_safe` of the registered tools should be True to make sure the tool calls are executed concurrently. To assert: 1. The events (by assert the dict generated by model_dump) 2. The final reply message 3. The agent state (Before and after replyStream) """ # Register concurrent tools conc_tool = MockConcurrentTool() self.agent.toolkit = Toolkit( tools=[conc_tool], ) # Create tool call IDs tool_call_id_1 = "tool_call_1" tool_call_id_2 = "tool_call_2" # Set up mock responses with tool calls self.model.set_responses( [ [ ChatResponse( content=[ ToolCallBlock( id=tool_call_id_1, name="mock_concurrent_tool", input='{"input": "test1"}', ), ToolCallBlock( id=tool_call_id_2, name="mock_concurrent_tool", input='{"input": "test2"}', ), ], is_last=True, usage=None, ), ], [ ChatResponse( content=[TextBlock(text="All done")], is_last=False, ), ChatResponse( content=[TextBlock(text="All done")], is_last=True, ), ], ], ) # Collect all events events = [] async for event in self.agent.reply_stream( UserMsg(name="user", content="Test"), ): events.append(event.model_dump()) # Verify events session_id = self.agent.state.session_id reply_id = self.agent.state.reply_id # For concurrent execution, the order of tool results may vary # Split events into: prefix (before concurrent), concurrent part, # suffix (after concurrent) # Expected prefix events (before tool execution) expected_prefix = [ { "type": "REPLY_START", "session_id": session_id, "name": "Friday", "role": "assistant", }, {"type": "MODEL_CALL_START", "model_name": "mock-model"}, { "type": "MODEL_CALL_END", "input_tokens": 0, "output_tokens": 0, "finished_reason": "completed", }, ] # Expected concurrent events (order may vary) expected_concurrent = [ { "type": "TOOL_RESULT_START", "tool_call_id": tool_call_id_1, "tool_call_name": "mock_concurrent_tool", }, { "type": "TOOL_RESULT_TEXT_DELTA", "tool_call_id": tool_call_id_1, "delta": "Concurrent result: test1", }, { "type": "TOOL_RESULT_END", "tool_call_id": tool_call_id_1, "state": "success", }, { "type": "TOOL_RESULT_START", "tool_call_id": tool_call_id_2, "tool_call_name": "mock_concurrent_tool", }, { "type": "TOOL_RESULT_TEXT_DELTA", "tool_call_id": tool_call_id_2, "delta": "Concurrent result: test2", }, { "type": "TOOL_RESULT_END", "tool_call_id": tool_call_id_2, "state": "success", }, ] # Expected suffix events (final model call with pure text) expected_suffix = [ {"type": "MODEL_CALL_START", "model_name": "mock-model"}, {"type": "TEXT_BLOCK_START", "block_id": AnyString()}, { "type": "TEXT_BLOCK_DELTA", "block_id": AnyString(), "delta": "All done", }, {"type": "TEXT_BLOCK_END", "block_id": AnyString()}, { "type": "MODEL_CALL_END", "input_tokens": 0, "output_tokens": 0, "finished_reason": "completed", }, { "type": "REPLY_END", "session_id": session_id, "finished_reason": "completed", }, ] # Assert prefix events (fixed order) basic_dict = self._get_event_base(reply_id) prefix_len = len(expected_prefix) self.assertListEqual( events[:prefix_len], [{**basic_dict, **_} for _ in expected_prefix], ) # Assert concurrent events (order may vary) concurrent_len = len(expected_concurrent) concurrent_events = events[prefix_len : prefix_len + concurrent_len] # Check length matches self.assertEqual(len(concurrent_events), len(expected_concurrent)) # Check each expected event is in the actual events for expected_event in expected_concurrent: expected_with_base = {**basic_dict, **expected_event} self.assertIn(expected_with_base, concurrent_events) # Assert suffix events (fixed order) suffix_events = events[prefix_len + concurrent_len :] self.assertListEqual( suffix_events, [{**basic_dict, **_} for _ in expected_suffix], ) # Assert context after reply_stream msg_base = self._get_msg_base() expected_context = [ { "name": "user", "role": "user", "content": [ { "type": "text", "id": AnyString(), "text": "Test", }, ], "finished_at": AnyString(), }, { "content": [ { "type": "tool_call", "id": tool_call_id_1, "name": "mock_concurrent_tool", "input": '{"input": "test1"}', "state": "finished", "suggested_rules": [], }, { "type": "tool_call", "id": tool_call_id_2, "name": "mock_concurrent_tool", "input": '{"input": "test2"}', "state": "finished", "suggested_rules": [], }, { "type": "tool_result", "id": AnyString(), "output": [ { "type": "text", "id": AnyString(), "text": "Concurrent result: test1", }, ], "name": "mock_concurrent_tool", "state": "success", "metadata": {}, }, { "type": "tool_result", "id": AnyString(), "output": [ { "type": "text", "id": AnyString(), "text": "Concurrent result: test2", }, ], "name": "mock_concurrent_tool", "state": "success", "metadata": {}, }, { "type": "text", "id": AnyString(), "text": "All done", }, ], }, ] 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) async def test_streaming_mixed_tool_calls(self) -> None: """Test the streaming model inference with both sequential and concurrent tool calls generated. Test only the replyStream interface. To assert: 1. The events (by assert the dict generated by model_dump) 2. The final reply message 3. The agent state (Before and after replyStream) """ # Register both sequential and concurrent tools seq_tool = MockSequentialTool() conc_tool = MockConcurrentTool() self.agent.toolkit = Toolkit( tools=[seq_tool, conc_tool], ) # Create tool call IDs tool_call_id_1 = "tool_call_1" tool_call_id_2 = "tool_call_2" tool_call_id_3 = "tool_call_3" # Set up mock responses with mixed tool calls self.model.set_responses( [ [ ChatResponse( content=[ ToolCallBlock( id=tool_call_id_1, name="mock_sequential_tool", input='{"input": "seq1"}', ), ToolCallBlock( id=tool_call_id_2, name="mock_concurrent_tool", input='{"input": "conc1"}', ), ToolCallBlock( id=tool_call_id_3, name="mock_concurrent_tool", input='{"input": "conc2"}', ), ], is_last=True, usage=None, ), ], [ ChatResponse( content=[TextBlock(text="All done")], is_last=False, ), ChatResponse( content=[TextBlock(text="All done")], is_last=True, ), ], ], ) # Collect all events events = [] async for event in self.agent.reply_stream( UserMsg(name="user", content="Test"), ): events.append(event.model_dump()) # Verify events session_id = self.agent.state.session_id reply_id = self.agent.state.reply_id # For mixed execution: sequential tool first, then concurrent tools # Split events into: prefix (before tool execution), sequential part, # concurrent part, suffix (final model call) # Expected prefix events (before tool execution) expected_prefix = [ { "type": "REPLY_START", "session_id": session_id, "name": "Friday", "role": "assistant", }, {"type": "MODEL_CALL_START", "model_name": "mock-model"}, { "type": "MODEL_CALL_END", "input_tokens": 0, "output_tokens": 0, "finished_reason": "completed", }, ] # Expected sequential tool execution events expected_sequential = [ { "type": "TOOL_RESULT_START", "tool_call_id": tool_call_id_1, "tool_call_name": "mock_sequential_tool", }, { "type": "TOOL_RESULT_TEXT_DELTA", "tool_call_id": tool_call_id_1, "delta": "Sequential result: seq1", }, { "type": "TOOL_RESULT_END", "tool_call_id": tool_call_id_1, "state": "success", "metadata": {}, }, ] # Expected concurrent events (order may vary) expected_concurrent: list[dict[str, Any]] = [ { "type": "TOOL_RESULT_START", "tool_call_id": tool_call_id_2, "tool_call_name": "mock_concurrent_tool", }, { "type": "TOOL_RESULT_TEXT_DELTA", "tool_call_id": tool_call_id_2, "delta": "Concurrent result: conc1", }, { "type": "TOOL_RESULT_END", "tool_call_id": tool_call_id_2, "state": "success", }, { "type": "TOOL_RESULT_START", "tool_call_id": tool_call_id_3, "tool_call_name": "mock_concurrent_tool", }, { "type": "TOOL_RESULT_TEXT_DELTA", "tool_call_id": tool_call_id_3, "delta": "Concurrent result: conc2", }, { "type": "TOOL_RESULT_END", "tool_call_id": tool_call_id_3, "state": "success", "metadata": {}, }, ] # Expected suffix events (final model call with pure text) expected_suffix = [ {"type": "MODEL_CALL_START", "model_name": "mock-model"}, {"type": "TEXT_BLOCK_START", "block_id": AnyString()}, { "type": "TEXT_BLOCK_DELTA", "block_id": AnyString(), "delta": "All done", }, {"type": "TEXT_BLOCK_END", "block_id": AnyString()}, { "type": "MODEL_CALL_END", "input_tokens": 0, "output_tokens": 0, "finished_reason": "completed", }, { "type": "REPLY_END", "session_id": session_id, "finished_reason": "completed", }, ] # Assert prefix events (fixed order) basic_dict = self._get_event_base(reply_id) prefix_len = len(expected_prefix) self.assertListEqual( events[:prefix_len], [{**basic_dict, **_} for _ in expected_prefix], ) # Assert sequential events (fixed order) sequential_len = len(expected_sequential) self.assertListEqual( events[prefix_len : prefix_len + sequential_len], [{**basic_dict, **_} for _ in expected_sequential], ) # Assert concurrent events (order may vary) concurrent_len = len(expected_concurrent) concurrent_events = events[ prefix_len + sequential_len : prefix_len + sequential_len + concurrent_len ] # Check length matches self.assertEqual(len(concurrent_events), len(expected_concurrent)) # Check each expected event is in the actual events for expected_event in expected_concurrent: expected_with_base = {**basic_dict, **expected_event} self.assertIn(expected_with_base, concurrent_events) # Assert suffix events (fixed order) suffix_events = events[prefix_len + sequential_len + concurrent_len :] self.assertListEqual( suffix_events, [{**basic_dict, **_} for _ in expected_suffix], ) # Assert context after reply_stream msg_base = self._get_msg_base() expected_context = [ { "name": "user", "role": "user", "content": [ { "type": "text", "id": AnyString(), "text": "Test", }, ], "finished_at": AnyString(), }, { "content": [ { "type": "tool_call", "id": tool_call_id_1, "name": "mock_sequential_tool", "input": '{"input": "seq1"}', "state": "finished", "suggested_rules": [], }, { "type": "tool_call", "id": tool_call_id_2, "name": "mock_concurrent_tool", "input": '{"input": "conc1"}', "state": "finished", "suggested_rules": [], }, { "type": "tool_call", "id": tool_call_id_3, "name": "mock_concurrent_tool", "input": '{"input": "conc2"}', "state": "finished", "suggested_rules": [], }, { "type": "tool_result", "id": AnyString(), "name": "mock_sequential_tool", "state": "success", "output": [ { "type": "text", "id": AnyString(), "text": "Sequential result: seq1", }, ], "metadata": {}, }, { "type": "tool_result", "id": AnyString(), "name": "mock_concurrent_tool", "state": "success", "output": [ { "type": "text", "id": AnyString(), "text": "Concurrent result: conc1", }, ], "metadata": {}, }, { "type": "tool_result", "id": AnyString(), "name": "mock_concurrent_tool", "state": "success", "output": [ { "type": "text", "id": AnyString(), "text": "Concurrent result: conc2", }, ], "metadata": {}, }, { "type": "text", "id": AnyString(), "text": "All done", }, ], }, ] 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) async def asyncTearDown(self) -> None: """The async teardown method."""