Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions strands-py/src/strands/multiagent/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,13 +28,15 @@ 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.
INTERRUPTED: Task was interrupted by user.
"""

PENDING = "pending"
DISPATCHED = "dispatched"
EXECUTING = "executing"
COMPLETED = "completed"
FAILED = "failed"
Expand Down
29 changes: 21 additions & 8 deletions strands-py/src/strands/multiagent/graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]

Expand Down Expand Up @@ -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.

Expand All @@ -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()
Expand Down
147 changes: 147 additions & 0 deletions strands-py/tests/strands/multiagent/test_graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Loading