Skip to content
Merged
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
Original file line number Diff line number Diff line change
Expand Up @@ -290,6 +290,9 @@ async def send_message(
f"Invalid selection result: {selection_results}. "
f"Expected selections to be a subset of valid target executor IDs: {self._target_ids}."
)
# A selection function can name a target more than once, for example when two tags on a
# message route to the same executor. Deliver to each selected target once.
selection_results = list(dict.fromkeys(selection_results))

if message.target_id:
# If the target ID is specified and the selection result contains it, send the message to that edge
Expand Down
55 changes: 55 additions & 0 deletions python/packages/core/tests/workflow/test_edge.py
Original file line number Diff line number Diff line change
Expand Up @@ -673,6 +673,61 @@ async def test_source_edge_group_with_selection_func_send_message_with_invalid_s
await edge_runner.send_message(message, state, ctx)


async def test_source_edge_group_with_selection_func_send_message_with_duplicate_selection() -> None:
"""Test that a target selected more than once receives the message once."""
source = MockExecutor(id="source_executor")
target1 = MockExecutor(id="target_executor_1")
target2 = MockExecutor(id="target_executor_2")

# e.g. a message tagged both "billing" and "refund" when both tags route to target1
edge_group = FanOutEdgeGroup(
source_id=source.id,
target_ids=[target1.id, target2.id],
selection_func=lambda data, target_ids: [target1.id, target1.id, target2.id],
)

executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2}
edge_runner = create_edge_runner(edge_group, executors)
state = State()
ctx = InProcRunnerContext()

data = MockMessage(data="test")
message = WorkflowMessage(data=data, source_id=source.id)

success = await edge_runner.send_message(message, state, ctx)

assert success is True
assert target1.call_count == 1
assert target2.call_count == 1


async def test_source_edge_group_with_selection_func_send_message_with_duplicate_selection_keeps_order() -> None:
"""Test that deduplicated targets are delivered in the order the selection function first names them."""
source = MockExecutor(id="source_executor")
target1 = MockExecutor(id="target_executor_1")
target2 = MockExecutor(id="target_executor_2")

edge_group = FanOutEdgeGroup(
source_id=source.id,
target_ids=[target1.id, target2.id],
selection_func=lambda data, target_ids: [target2.id, target1.id, target2.id],
)

executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2}
edge_runner = create_edge_runner(edge_group, executors)
state = State()
ctx = InProcRunnerContext()

data = MockMessage(data="test")
message = WorkflowMessage(data=data, source_id=source.id)

with patch("agent_framework._workflows._edge_runner.EdgeRunner._execute_on_target") as mock_send:
success = await edge_runner.send_message(message, state, ctx)

assert success is True
assert [call.args[0] for call in mock_send.call_args_list] == [target2.id, target1.id]


async def test_source_edge_group_with_selection_func_send_message_with_target() -> None:
"""Test sending a message through a fan-out group with a selection func with a target."""
source = MockExecutor(id="source_executor")
Expand Down
Loading