Merge branch 'main' of https://github.com/strands-agents/sdk-python into avoid-event-loop-block

Author: pgrayy

src/strands/event_loop/event_loop.py

@@ -132,14 +132,14 @@ async def event_loop_cycle(agent: "Agent", invocation_state: dict[str, Any]) ->
             model_id=model_id,
         )
         with trace_api.use_span(model_invoke_span):
-            tool_specs = agent.tool_registry.get_all_tool_specs()
-
             agent.hooks.invoke_callbacks(
                 BeforeModelInvocationEvent(
                     agent=agent,
                 )
             )
 
+            tool_specs = agent.tool_registry.get_all_tool_specs()
+
             try:
                 async for event in stream_messages(agent.model, agent.system_prompt, agent.messages, tool_specs):
                     if not isinstance(event, ModelStopReason):

src/strands/event_loop/streaming.py

@@ -289,7 +289,6 @@ async def process_stream(chunks: AsyncIterable[StreamEvent]) -> AsyncGenerator[T
         "text": "",
         "current_tool_use": {},
         "reasoningText": "",
-        "signature": "",
         "citationsContent": [],
     }
     state["content"] = state["message"]["content"]

src/strands/multiagent/graph.py

@@ -469,41 +469,32 @@ async def _execute_graph(self) -> None:
             ready_nodes.clear()
 
             # Execute current batch of ready nodes concurrently
-            tasks = [
-                asyncio.create_task(self._execute_node(node))
-                for node in current_batch
-                if node not in self.state.completed_nodes
-            ]
+            tasks = [asyncio.create_task(self._execute_node(node)) for node in current_batch]
 
             for task in tasks:
                 await task
 
             # Find newly ready nodes after batch execution
-            ready_nodes.extend(self._find_newly_ready_nodes())
+            # We add all nodes in current batch as completed batch,
+            # because a failure would throw exception and code would not make it here
+            ready_nodes.extend(self._find_newly_ready_nodes(current_batch))
 
-    def _find_newly_ready_nodes(self) -> list["GraphNode"]:
+    def _find_newly_ready_nodes(self, completed_batch: list["GraphNode"]) -> list["GraphNode"]:
         """Find nodes that became ready after the last execution."""
         newly_ready = []
         for _node_id, node in self.nodes.items():
-            if (
-                node not in self.state.completed_nodes
-                and node not in self.state.failed_nodes
-                and self._is_node_ready_with_conditions(node)
-            ):
+            if self._is_node_ready_with_conditions(node, completed_batch):
                 newly_ready.append(node)
         return newly_ready
 
-    def _is_node_ready_with_conditions(self, node: GraphNode) -> bool:
+    def _is_node_ready_with_conditions(self, node: GraphNode, completed_batch: list["GraphNode"]) -> bool:
         """Check if a node is ready considering conditional edges."""
         # Get incoming edges to this node
         incoming_edges = [edge for edge in self.edges if edge.to_node == node]
 
-        if not incoming_edges:
-            return node in self.entry_points
-
         # Check if at least one incoming edge condition is satisfied
         for edge in incoming_edges:
-            if edge.from_node in self.state.completed_nodes:
+            if edge.from_node in completed_batch:
                 if edge.should_traverse(self.state):
                     logger.debug(
                         "from=<%s>, to=<%s> | edge ready via satisfied condition", edge.from_node.node_id, node.node_id

tests/strands/event_loop/test_streaming.py

@@ -1,10 +1,12 @@
 import unittest.mock
+from typing import cast
 
 import pytest
 
 import strands
 import strands.event_loop
-from strands.types._events import TypedEvent
+from strands.types._events import ModelStopReason, TypedEvent
+from strands.types.content import Message
 from strands.types.streaming import (
     ContentBlockDeltaEvent,
     ContentBlockStartEvent,
@@ -565,6 +567,88 @@ async def test_process_stream(response, exp_events, agenerator, alist):
     assert non_typed_events == []
 
 
+def _get_message_from_event(event: ModelStopReason) -> Message:
+    return cast(Message, event["stop"][1])
+
+
+@pytest.mark.asyncio
+async def test_process_stream_with_no_signature(agenerator, alist):
+    response = [
+        {"messageStart": {"role": "assistant"}},
+        {
+            "contentBlockDelta": {
+                "delta": {"reasoningContent": {"text": 'User asks: "Reason about 2+2" so I will do that'}},
+                "contentBlockIndex": 0,
+            }
+        },
+        {"contentBlockDelta": {"delta": {"reasoningContent": {"text": "."}}, "contentBlockIndex": 0}},
+        {"contentBlockStop": {"contentBlockIndex": 0}},
+        {
+            "contentBlockDelta": {
+                "delta": {"text": "Sure! Let’s do it"},
+                "contentBlockIndex": 1,
+            }
+        },
+        {"contentBlockStop": {"contentBlockIndex": 1}},
+        {"messageStop": {"stopReason": "end_turn"}},
+        {
+            "metadata": {
+                "usage": {"inputTokens": 112, "outputTokens": 764, "totalTokens": 876},
+                "metrics": {"latencyMs": 2970},
+            }
+        },
+    ]
+
+    stream = strands.event_loop.streaming.process_stream(agenerator(response))
+
+    last_event = cast(ModelStopReason, (await alist(stream))[-1])
+
+    message = _get_message_from_event(last_event)
+
+    assert "signature" not in message["content"][0]["reasoningContent"]["reasoningText"]
+    assert message["content"][1]["text"] == "Sure! Let’s do it"
+
+
+@pytest.mark.asyncio
+async def test_process_stream_with_signature(agenerator, alist):
+    response = [
+        {"messageStart": {"role": "assistant"}},
+        {
+            "contentBlockDelta": {
+                "delta": {"reasoningContent": {"text": 'User asks: "Reason about 2+2" so I will do that'}},
+                "contentBlockIndex": 0,
+            }
+        },
+        {"contentBlockDelta": {"delta": {"reasoningContent": {"text": "."}}, "contentBlockIndex": 0}},
+        {"contentBlockDelta": {"delta": {"reasoningContent": {"signature": "test-"}}, "contentBlockIndex": 0}},
+        {"contentBlockDelta": {"delta": {"reasoningContent": {"signature": "signature"}}, "contentBlockIndex": 0}},
+        {"contentBlockStop": {"contentBlockIndex": 0}},
+        {
+            "contentBlockDelta": {
+                "delta": {"text": "Sure! Let’s do it"},
+                "contentBlockIndex": 1,
+            }
+        },
+        {"contentBlockStop": {"contentBlockIndex": 1}},
+        {"messageStop": {"stopReason": "end_turn"}},
+        {
+            "metadata": {
+                "usage": {"inputTokens": 112, "outputTokens": 764, "totalTokens": 876},
+                "metrics": {"latencyMs": 2970},
+            }
+        },
+    ]
+
+    stream = strands.event_loop.streaming.process_stream(agenerator(response))
+
+    last_event = cast(ModelStopReason, (await alist(stream))[-1])
+
+    message = _get_message_from_event(last_event)
+
+    assert message["content"][0]["reasoningContent"]["reasoningText"]["signature"] == "test-signature"
+    assert message["content"][1]["text"] == "Sure! Let’s do it"
+
+
 @pytest.mark.asyncio
 async def test_stream_messages(agenerator, alist):
     mock_model = unittest.mock.MagicMock()