bidi - fix record direct tool call (#1300)

Author: pgrayy

src/strands/agent/agent.py

@@ -624,8 +624,7 @@ async def _run_loop(
         try:
             yield InitEventLoopEvent()
 
-            for message in messages:
-                await self._append_message(message)
+            await self._append_messages(*messages)
 
             structured_output_context = StructuredOutputContext(
                 structured_output_model or self._default_structured_output_model
@@ -715,7 +714,7 @@ async def _convert_prompt_to_messages(self, prompt: AgentInput) -> Messages:
                 tool_use_ids = [
                     content["toolUse"]["toolUseId"] for content in self.messages[-1]["content"] if "toolUse" in content
                 ]
-                await self._append_message(
+                await self._append_messages(
                     {
                         "role": "user",
                         "content": generate_missing_tool_result_content(tool_use_ids),
@@ -811,10 +810,11 @@ def _initialize_system_prompt(
         else:
             return None, None
 
-    async def _append_message(self, message: Message) -> None:
-        """Appends a message to the agent's list of messages and invokes the callbacks for the MessageCreatedEvent."""
-        self.messages.append(message)
-        await self.hooks.invoke_callbacks_async(MessageAddedEvent(agent=self, message=message))
+    async def _append_messages(self, *messages: Message) -> None:
+        """Appends messages to history and invoke the callbacks for the MessageAddedEvent."""
+        for message in messages:
+            self.messages.append(message)
+            await self.hooks.invoke_callbacks_async(MessageAddedEvent(agent=self, message=message))
 
     def _redact_user_content(self, content: list[ContentBlock], redact_message: str) -> list[ContentBlock]:
         """Redact user content preserving toolResult blocks.

src/strands/event_loop/event_loop.py

@@ -230,7 +230,7 @@ async def event_loop_cycle(
             )
         structured_output_context.set_forced_mode()
         logger.debug("Forcing structured output tool")
-        await agent._append_message(
+        await agent._append_messages(
             {"role": "user", "content": [{"text": "You must format the previous response as structured output."}]}
         )
 

src/strands/experimental/bidi/agent/agent.py

@@ -26,9 +26,9 @@
 from ....tools.executors._executor import ToolExecutor
 from ....tools.registry import ToolRegistry
 from ....tools.watcher import ToolWatcher
-from ....types.content import Messages
+from ....types.content import Message, Messages
 from ....types.tools import AgentTool
-from ...hooks.events import BidiAgentInitializedEvent
+from ...hooks.events import BidiAgentInitializedEvent, BidiMessageAddedEvent
 from ...tools import ToolProvider
 from .._async import stop_all
 from ..models.model import BidiModel
@@ -167,6 +167,9 @@ def __init__(
         # TODO: Determine if full support is required
         self._interrupt_state = _InterruptState()
 
+        # Lock to ensure that paired messages are added to history in sequence without interference
+        self._message_lock = asyncio.Lock()
+
         self._started = False
 
     @property
@@ -396,3 +399,17 @@ async def run_outputs(inputs_task: asyncio.Task) -> None:
             output_stops = [output.stop for output in outputs if isinstance(output, BidiOutput)]
 
             await stop_all(*input_stops, *output_stops, self.stop)
+
+    async def _append_messages(self, *messages: Message) -> None:
+        """Append messages to history in sequence without interference.
+
+        The message lock ensures that paired messages are added to history in sequence without interference. For
+        example, tool use and tool result messages must be added adjacent to each other.
+
+        Args:
+            *messages: List of messages to add into history.
+        """
+        async with self._message_lock:
+            for message in messages:
+                self.messages.append(message)
+                await self.hooks.invoke_callbacks_async(BidiMessageAddedEvent(agent=self, message=message))

src/strands/experimental/bidi/agent/loop.py

@@ -15,7 +15,6 @@
     BidiAfterInvocationEvent,
     BidiBeforeConnectionRestartEvent,
     BidiBeforeInvocationEvent,
-    BidiMessageAddedEvent,
 )
 from ...hooks.events import (
     BidiInterruptionEvent as BidiInterruptionHookEvent,
@@ -51,8 +50,6 @@ class _BidiAgentLoop:
             that tools can access via their invocation_state parameter.
         _send_gate: Gate the sending of events to the model.
             Blocks when agent is reseting the model connection after timeout.
-        _message_lock: Lock to ensure that paired messages are added to history in sequence without interference.
-            For example, tool use and tool result messages must be added adjacent to each other.
     """
 
     def __init__(self, agent: "BidiAgent") -> None:
@@ -70,7 +67,6 @@ def __init__(self, agent: "BidiAgent") -> None:
         self._invocation_state: dict[str, Any]
 
         self._send_gate = asyncio.Event()
-        self._message_lock = asyncio.Lock()
 
     async def start(self, invocation_state: dict[str, Any] | None = None) -> None:
         """Start the agent loop.
@@ -145,7 +141,7 @@ async def send(self, event: BidiInputEvent | ToolResultEvent) -> None:
 
         if isinstance(event, BidiTextInputEvent):
             message: Message = {"role": "user", "content": [{"text": event.text}]}
-            await self._add_messages(message)
+            await self._agent._append_messages(message)
 
         await self._agent.model.send(event)
 
@@ -224,7 +220,7 @@ async def _run_model(self) -> None:
                 if isinstance(event, BidiTranscriptStreamEvent):
                     if event["is_final"]:
                         message: Message = {"role": event["role"], "content": [{"text": event["text"]}]}
-                        await self._add_messages(message)
+                        await self._agent._append_messages(message)
 
                 elif isinstance(event, ToolUseStreamEvent):
                     tool_use = event["current_tool_use"]
@@ -282,7 +278,7 @@ async def _run_tool(self, tool_use: ToolUse) -> None:
 
             tool_use_message: Message = {"role": "assistant", "content": [{"toolUse": tool_use}]}
             tool_result_message: Message = {"role": "user", "content": [{"toolResult": tool_result_event.tool_result}]}
-            await self._add_messages(tool_use_message, tool_result_message)
+            await self._agent._append_messages(tool_use_message, tool_result_message)
 
             await self._event_queue.put(ToolResultMessageEvent(tool_result_message))
 
@@ -300,16 +296,3 @@ async def _run_tool(self, tool_use: ToolUse) -> None:
 
         except Exception as error:
             await self._event_queue.put(error)
-
-    async def _add_messages(self, *messages: Message) -> None:
-        """Add messages to history in sequence without interference.
-
-        Args:
-            *messages: List of messages to add into history.
-        """
-        async with self._message_lock:
-            for message in messages:
-                self._agent.messages.append(message)
-                await self._agent.hooks.invoke_callbacks_async(
-                    BidiMessageAddedEvent(agent=self._agent, message=message)
-                )

src/strands/tools/_caller.py

@@ -195,10 +195,7 @@ async def _record_tool_execution(
         }
 
         # Add to message history
-        await self._agent._append_message(user_msg)
-        await self._agent._append_message(tool_use_msg)
-        await self._agent._append_message(tool_result_msg)
-        await self._agent._append_message(assistant_msg)
+        await self._agent._append_messages(user_msg, tool_use_msg, tool_result_msg, assistant_msg)
 
     def _filter_tool_parameters_for_recording(self, tool_name: str, input_params: dict[str, Any]) -> dict[str, Any]:
         """Filter input parameters to only include those defined in the tool specification.

tests/strands/event_loop/test_event_loop_structured_output.py

@@ -42,7 +42,7 @@ def mock_agent():
     agent.trace_span = None
     agent.trace_attributes = {}
     agent.tool_executor = Mock()
-    agent._append_message = AsyncMock()
+    agent._append_messages = AsyncMock()
 
     # Set up _interrupt_state properly
     agent._interrupt_state = Mock()
@@ -186,8 +186,8 @@ async def test_event_loop_forces_structured_output_on_end_turn(
         await alist(stream)
 
         # Should have appended a message to force structured output
-        mock_agent._append_message.assert_called_once()
-        args = mock_agent._append_message.call_args[0][0]
+        mock_agent._append_messages.assert_called_once()
+        args = mock_agent._append_messages.call_args[0][0]
         assert args["role"] == "user"
 
         # Should have called recurse_event_loop with the context

tests/strands/experimental/bidi/agent/test_loop.py

@@ -4,12 +4,10 @@
 import pytest_asyncio
 
 from strands import tool
+from strands.experimental.bidi import BidiAgent
 from strands.experimental.bidi.agent.loop import _BidiAgentLoop
 from strands.experimental.bidi.models import BidiModelTimeoutError
 from strands.experimental.bidi.types.events import BidiConnectionRestartEvent, BidiTextInputEvent
-from strands.hooks import HookRegistry
-from strands.tools.executors import SequentialToolExecutor
-from strands.tools.registry import ToolRegistry
 from strands.types._events import ToolResultEvent, ToolResultMessageEvent, ToolUseStreamEvent
 
 
@@ -24,20 +22,12 @@ async def func():
 
 @pytest.fixture
 def agent(time_tool):
-    mock = unittest.mock.Mock()
-    mock.hooks = HookRegistry()
-    mock.messages = []
-    mock.model = unittest.mock.AsyncMock()
-    mock.tool_executor = SequentialToolExecutor()
-    mock.tool_registry = ToolRegistry()
-    mock.tool_registry.process_tools([time_tool])
-    
-    return mock
+    return BidiAgent(model=unittest.mock.AsyncMock(), tools=[time_tool])
 
 
 @pytest_asyncio.fixture
 async def loop(agent):
-    return _BidiAgentLoop(agent)
+    return agent._loop
 
 
 @pytest.mark.asyncio

tests_integ/bidi/tools/__init__.py

@@ -0,0 +1,74 @@
+import unittest.mock
+
+import pytest
+
+from strands import tool
+from strands.experimental.bidi.agent import BidiAgent
+
+
+@pytest.fixture
+def weather_tool():
+    @tool(name="weather_tool")
+    def func(city_name: str) -> str:
+        return f"city_name=<{city_name}> | sunny"
+
+    return func
+
+
+@pytest.fixture
+def agent(weather_tool):
+    return BidiAgent(record_direct_tool_call=True, tools=[weather_tool])
+
+
+def test_bidi_agent_tool_direct_call(agent):
+    tru_result = agent.tool.weather_tool(city_name="new york")
+    exp_result = {
+        "content": [{"text": "city_name=<new york> | sunny"}],
+        "status": "success",
+        "toolUseId": unittest.mock.ANY,
+    }
+    assert tru_result == exp_result
+    
+    tru_messages = agent.messages
+    exp_messages = [
+        {
+            "content": [
+                {
+                    "text": (
+                        "agent.tool.weather_tool direct tool call.\n"
+                        'Input parameters: {"city_name": "new york"}\n'
+                    ),
+                },
+            ],
+            "role": "user",
+        },
+        {
+            "content": [
+                {
+                    "toolUse": {
+                        "input": {"city_name": "new york"},
+                        "name": "weather_tool",
+                        "toolUseId": unittest.mock.ANY,
+                    },
+                },
+            ],
+            "role": "assistant",
+        },
+        {
+            "content": [
+                {
+                    "toolResult": {
+                        "content": [{"text": "city_name=<new york> | sunny"}],
+                        "status": "success",
+                        "toolUseId": unittest.mock.ANY,
+                    },
+                },
+            ],
+            "role": "user",
+        },
+        {
+            "content": [{"text": "agent.tool.weather_tool was called."}],
+            "role": "assistant",
+        },
+    ]
+    assert tru_messages == exp_messages