diff --git a/strands-py/src/strands/multiagent/base.py b/strands-py/src/strands/multiagent/base.py index e7a3e3c9c3..38f83a7974 100644 --- a/strands-py/src/strands/multiagent/base.py +++ b/strands-py/src/strands/multiagent/base.py @@ -28,6 +28,7 @@ class Status(Enum): Attributes: PENDING: Task has not started execution yet. + DISPATCHED: Node has been dispatched but execution has not started yet. EXECUTING: Task is currently running. COMPLETED: Task finished successfully. FAILED: Task encountered an error and could not complete. @@ -35,6 +36,7 @@ class Status(Enum): """ PENDING = "pending" + DISPATCHED = "dispatched" EXECUTING = "executing" COMPLETED = "completed" FAILED = "failed" diff --git a/strands-py/src/strands/multiagent/graph.py b/strands-py/src/strands/multiagent/graph.py index 0a2378f73c..2dd8476bef 100644 --- a/strands-py/src/strands/multiagent/graph.py +++ b/strands-py/src/strands/multiagent/graph.py @@ -829,6 +829,13 @@ async def _execute_nodes_parallel( event_queue: asyncio.Queue[Any | None | Exception] = asyncio.Queue() + # Mark each node as DISPATCHED before scheduling its task. _execute_node flips the + # status to EXECUTING once the node's generator body starts; this marks the window in + # between so a serialize captured there (node dispatched but not yet executing) treats + # the node as in-flight rather than a bypassed dead branch — see _compute_pending_sources. + for node in nodes: + node.execution_status = Status.DISPATCHED + # Start all node streams as independent tasks tasks = [asyncio.create_task(self._stream_node_to_queue(node, event_queue, invocation_state)) for node in nodes] @@ -1349,11 +1356,15 @@ def _compute_pending_sources(self, completed_nodes: set[GraphNode]) -> set[Graph branch*, so a fan-in node stays not-ready while any parent that will still produce a result is outstanding. It has two seeds: - - In-flight nodes (``execution_status == EXECUTING``): serialize runs from the per-node - ``AfterNodeCallEvent`` while sibling nodes in the same parallel batch are still executing. - Their partial progress is not persisted, so they re-run on resume. This is the - authoritative in-flight signal, and it catches siblings with no completed ancestor (e.g. - two concurrent entry points feeding one join) that forward reachability alone would miss. + - In-flight nodes (``execution_status in (EXECUTING, DISPATCHED)``): serialize runs + from the per-node ``AfterNodeCallEvent`` while sibling nodes in the same parallel batch + are still executing (EXECUTING), or while a node has been dispatched but its generator + body has not yet started (DISPATCHED). Their partial progress is not persisted, so they + re-run on resume. This is the authoritative in-flight signal, and it catches siblings + with no completed ancestor (e.g. two concurrent entry points feeding one join) that + forward reachability alone would miss. DISPATCHED specifically covers a node that was + scheduled but crashed before reaching EXECUTING — without it the node stays PENDING and + is treated as a bypassed dead branch, double-executing its fan-in children on resume. - Completed nodes: their traversable outgoing edges lead to nodes that will be scheduled next on resume. @@ -1373,9 +1384,11 @@ def _compute_pending_sources(self, completed_nodes: set[GraphNode]) -> set[Graph The set of not-yet-completed nodes that will execute on resume — in-flight nodes plus everything reachable from the completed or in-flight frontier via traversable edges. """ - executing_nodes = {node for node in self.nodes.values() if node.execution_status == Status.EXECUTING} - pending: set[GraphNode] = set(executing_nodes) - visited = completed_nodes | executing_nodes + in_flight_nodes = { + node for node in self.nodes.values() if node.execution_status in (Status.EXECUTING, Status.DISPATCHED) + } + pending: set[GraphNode] = set(in_flight_nodes) + visited = completed_nodes | in_flight_nodes frontier = list(visited) while frontier: source = frontier.pop() diff --git a/strands-py/tests/strands/multiagent/test_graph.py b/strands-py/tests/strands/multiagent/test_graph.py index d709224bba..308b3cbef2 100644 --- a/strands-py/tests/strands/multiagent/test_graph.py +++ b/strands-py/tests/strands/multiagent/test_graph.py @@ -3468,3 +3468,150 @@ def capture_on_left(event): join_input_text = " ".join(str(block) for block in join_input) assert "From left:" in join_input_text assert "From right:" in join_input_text + + +class TestResumeDispatchedSibling: + """Verify a fan-in node is not marked ready while a dispatched sibling has not yet started. + + A node that was dispatched (``asyncio.create_task``) but crashed before its generator body set + ``Status.EXECUTING`` is left ``DISPATCHED`` rather than reverting to ``PENDING``. Without the + DISPATCHED status such a node would be PENDING and — when it has no completed ancestor (e.g. + two concurrent entry points feeding one join) — ``_compute_pending_sources`` would not seed it, + so it would be treated as a bypassed dead branch: its edge into the fan-in would be dropped and + the fan-in marked ready off the already-completed sibling alone, double-executing it on resume. + Seeding DISPATCHED as in-flight keeps the fan-in not-ready until the dispatched sibling runs. + """ + + def test_dispatched_sibling_excluded_from_resume_frontier(self): + """join is excluded and right included when right is dispatched but not yet executing. + + Focused unit test that injects the DISPATCHED state directly + (``right.execution_status = DISPATCHED``) rather than driving the real dispatch path; + it pins ``_compute_ready_nodes_for_resume`` for the DISPATCHED window. The companion + ``test_resume_runs_fan_in_once_for_dispatched_sibling`` exercises the real dispatch + path end-to-end. On master (no DISPATCHED enum) this raises ``AttributeError`` on + ``Status.DISPATCHED`` — a valid red signal, distinct from the assertion failure of + the end-to-end test. + """ + builder = GraphBuilder() + builder.add_node(create_mock_agent("left"), "left") + builder.add_node(create_mock_agent("right"), "right") + builder.add_node(create_mock_agent("join"), "join") + builder.set_entry_point("left") + builder.set_entry_point("right") + builder.add_edge("left", "join") + builder.add_edge("right", "join") + graph = builder.build() + + node_left = graph.nodes["left"] + node_right = graph.nodes["right"] + + # left completed; right dispatched but not yet executing (crashed before EXECUTING), in + # none of the terminal sets and with no completed ancestor to reach it from forward + # reachability alone. + node_left.execution_status = Status.COMPLETED + node_right.execution_status = Status.DISPATCHED + + graph.state.status = Status.EXECUTING + graph.state.completed_nodes = {node_left} + + payload = graph.serialize_state() + + assert "join" not in payload["next_nodes_to_execute"], ( + f"join must not be ready while dispatched sibling right is outstanding, " + f"got {payload['next_nodes_to_execute']}" + ) + assert "right" in payload["next_nodes_to_execute"] + + ready_ids = {node.node_id for node in graph._compute_ready_nodes_for_resume()} + assert ready_ids == {"right"} + + @pytest.mark.asyncio + async def test_resume_runs_fan_in_once_for_dispatched_sibling(self): + """A restore from a dispatched-but-not-started snapshot runs the fan-in exactly once. + + Reproduces the real dispatch path (no state injection): left and right are concurrent + entry points feeding join. The production ``_execute_nodes_parallel`` path sets right to + DISPATCHED and schedules its task; a thin wrapper holds right in that window — before + ``_execute_node`` flips it to EXECUTING — so the snapshot captured from left's + ``AfterNodeCallEvent`` sees right genuinely DISPATCHED. The snapshot is resumed into a + fresh graph; join must run exactly once with both parents' outputs present. + + On master (no DISPATCHED status) right stays PENDING, so it is treated as a bypassed + dead branch: join is marked ready off left alone and lands in + ``next_nodes_to_execute``, which double-executes it on resume — a real assertion + failure, not an AttributeError. + """ + release_right = asyncio.Event() + + agent_left = create_mock_agent("left", "left done") + agent_right = create_mock_agent("right", "right done") + + builder = GraphBuilder() + builder.add_node(agent_left, "left") + builder.add_node(agent_right, "right") + builder.add_node(create_mock_agent("join", "join done"), "join") + builder.set_entry_point("left") + builder.set_entry_point("right") + builder.add_edge("left", "join") + builder.add_edge("right", "join") + graph = builder.build() + + # Capture the real implementation before patching so the wrapper can delegate to it. + real_stream_node_to_queue = graph._stream_node_to_queue + + async def blocking_right_dispatch(node, event_queue, invocation_state): + # Hold right in the DISPATCHED window: its task has been created (DISPATCHED set by + # _execute_nodes_parallel) but _execute_node has not run yet, so EXECUTING is never + # set until release. This widens the one-tick create_task→_execute_node gap into a + # deterministic capture point reached by left's AfterNodeCallEvent. + if node.node_id == "right": + await release_right.wait() + await real_stream_node_to_queue(node, event_queue, invocation_state) + + graph._stream_node_to_queue = blocking_right_dispatch # type: ignore[method-assign] + + captured: dict = {} + + def capture_on_left(event): + # left completes while right is still DISPATCHED — snapshot the live-sibling state, + # then release right so the first graph can finish. + if event.node_id == "left" and "payload" not in captured: + captured["payload"] = event.source.serialize_state() + release_right.set() + + graph.add_hook(capture_on_left, AfterNodeCallEvent) + + await graph.invoke_async("two entry points") + + assert "payload" in captured, "left's AfterNodeCallEvent should fire while right is DISPATCHED" + assert "right" in captured["payload"]["next_nodes_to_execute"] + assert "join" not in captured["payload"]["next_nodes_to_execute"] + + # Resume into a fresh graph with independent executor mocks, since the call-count and + # input assertions target the resumed graph's nodes. + builder2 = GraphBuilder() + builder2.add_node(create_mock_agent("left", "left done"), "left") + builder2.add_node(create_mock_agent("right", "right done"), "right") + agent_join2 = create_mock_agent("join", "join done") + builder2.add_node(agent_join2, "join") + builder2.set_entry_point("left") + builder2.set_entry_point("right") + builder2.add_edge("left", "join") + builder2.add_edge("right", "join") + graph2 = builder2.build() + + graph2.deserialize_state(captured["payload"]) + result = await graph2.invoke_async("two entry points") + + # join runs exactly once — a stale frontier including join would double-run it. + assert agent_join2.stream_async.call_count == 1 + assert result.status == Status.COMPLETED + assert "join" in result.results + + # join's single run sees both parents' outputs (left restored from the snapshot, right + # re-run on resume). + join_input = agent_join2.stream_async.call_args.args[0] + join_input_text = " ".join(str(block) for block in join_input) + assert "From left:" in join_input_text + assert "From right:" in join_input_text