From b0b8bd99e05f7e940b87e070d7cb591c2eb17a4e Mon Sep 17 00:00:00 2001 From: weed33834 Date: Fri, 7 Aug 2026 04:10:16 +0000 Subject: [PATCH 1/2] fix(workflows): preserve all trace contexts in FanInEdgeRunner aggregation FanInEdgeRunner collected trace contexts and source span IDs using the singular backward-compat properties (msg.trace_context / msg.source_span_id), which return only the first element of the plural lists. When a message arriving at a fan-in already carries multiple trace contexts (e.g. from a prior fan-in aggregation), all but the first were silently dropped. Iterate over the plural fields (trace_contexts / source_span_ids) and extend the aggregated lists so every trace context and source span ID from every source message is preserved. This keeps distributed tracing links intact for nested fan-in topologies. Added test_fan_in_preserves_multiple_trace_contexts_per_message that sends a message with two trace contexts through a fan-in and asserts all three contexts (2 + 1) reach the target executor. --- .../_workflows/_edge_runner.py | 16 ++- .../packages/core/tests/workflow/test_edge.py | 99 +++++++++++++++++++ 2 files changed, 112 insertions(+), 3 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_edge_runner.py b/python/packages/core/agent_framework/_workflows/_edge_runner.py index c26eca3c552..8955335e5b4 100644 --- a/python/packages/core/agent_framework/_workflows/_edge_runner.py +++ b/python/packages/core/agent_framework/_workflows/_edge_runner.py @@ -353,9 +353,19 @@ async def send_message( # Send aggregated data to target aggregated_data = [msg.data for msg in messages_to_send] - # Collect all trace contexts and source span IDs for fan-in linking - trace_contexts = [msg.trace_context for msg in messages_to_send if msg.trace_context] - source_span_ids = [msg.source_span_id for msg in messages_to_send if msg.source_span_id] + # Collect all trace contexts and source span IDs for fan-in linking. + # Iterate over the plural fields (trace_contexts / source_span_ids) + # so that messages carrying multiple contexts from a previous + # fan-in aggregation are fully preserved. Using the singular + # backward-compat properties would silently drop all but the + # first context per message. + trace_contexts: list[dict[str, str]] = [] + source_span_ids: list[str] = [] + for msg in messages_to_send: + if msg.trace_contexts: + trace_contexts.extend(msg.trace_contexts) + if msg.source_span_ids: + source_span_ids.extend(msg.source_span_ids) # Create a new Message object for the aggregated data aggregated_message = WorkflowMessage( diff --git a/python/packages/core/tests/workflow/test_edge.py b/python/packages/core/tests/workflow/test_edge.py index 422d530631e..dcb243a27c3 100644 --- a/python/packages/core/tests/workflow/test_edge.py +++ b/python/packages/core/tests/workflow/test_edge.py @@ -1191,6 +1191,105 @@ async def test_fan_in_edge_group_with_multiple_message_types_failed() -> None: ) +class TraceCapturingAggregator(Executor): + """Fan-in aggregator that captures the trace contexts passed to execute().""" + + def __init__(self, *, id: str) -> None: + super().__init__(id=id) + self.captured_trace_contexts: list[dict[str, str]] | None = None + self.captured_source_span_ids: list[str] | None = None + self.call_count: int = 0 + + @handler + async def mock_aggregator_handler(self, message: list[MockMessage], ctx: WorkflowContext) -> None: + self.call_count += 1 + + async def execute( + self, + message: WorkflowMessage, + source_executor_ids: list[str], + state: Any, + ctx: Any, + *, + trace_contexts: list[dict[str, str]] | None = None, + source_span_ids: list[str] | None = None, + ) -> None: + self.captured_trace_contexts = trace_contexts + self.captured_source_span_ids = source_span_ids + await super().execute( + message, + source_executor_ids, + state, + ctx, + trace_contexts=trace_contexts, + source_span_ids=source_span_ids, + ) + + +async def test_fan_in_preserves_multiple_trace_contexts_per_message() -> None: + """Fan-in must aggregate ALL trace contexts, not just the first per message. + + Each incoming message may carry multiple trace_contexts (e.g. when it is + itself the product of a prior fan-in). The aggregated message must include + every trace context and source span ID from every source message; using the + singular backward-compat properties (trace_context / source_span_id) would + silently drop all but the first context per message. + """ + source1 = MockExecutor(id="source_executor_1") + source2 = MockExecutor(id="source_executor_2") + target = TraceCapturingAggregator(id="target_executor") + + executors: dict[str, Executor] = {source1.id: source1, source2.id: source2, target.id: target} + edge_group = FanInEdgeGroup(source_ids=[source1.id, source2.id], target_id=target.id) + edge_runner = create_edge_runner(edge_group, executors) + + state = State() + ctx = InProcRunnerContext() + data = MockMessage(data="test") + + # Source 1 carries TWO trace contexts (simulating a prior fan-in aggregation) + multi_contexts_1 = [ + {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"}, + {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-aaaaaaaaaaaaaaaa-01"}, + ] + multi_span_ids_1 = ["00f067aa0ba902b7", "aaaaaaaaaaaaaaaa"] + + # Source 2 carries a single trace context + single_contexts_2 = [ + {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b8-01"}, + ] + single_span_ids_2 = ["00f067aa0ba902b8"] + + # Send first message (buffered) + assert await edge_runner.send_message( + WorkflowMessage( + data=data, source_id=source1.id, trace_contexts=multi_contexts_1, source_span_ids=multi_span_ids_1 + ), + state, + ctx, + ) + + # Send second message (triggers delivery) + assert await edge_runner.send_message( + WorkflowMessage( + data=data, source_id=source2.id, trace_contexts=single_contexts_2, source_span_ids=single_span_ids_2 + ), + state, + ctx, + ) + + # The target executor should have received ALL 3 trace contexts (2 + 1), + # not just 2 (one per message via the singular property). + assert target.call_count == 1 + assert target.captured_trace_contexts is not None + assert len(target.captured_trace_contexts) == 3 + assert target.captured_trace_contexts == multi_contexts_1 + single_contexts_2 + + assert target.captured_source_span_ids is not None + assert len(target.captured_source_span_ids) == 3 + assert target.captured_source_span_ids == multi_span_ids_1 + single_span_ids_2 + + # endregion FanInEdgeGroup # region SwitchCaseEdgeGroup From 6fdc8b5869d8e5e9edd32d5732e5eef77d60cdc1 Mon Sep 17 00:00:00 2001 From: weed33834 Date: Sat, 8 Aug 2026 04:47:10 +0000 Subject: [PATCH 2/2] fix: address Copilot review comments on trace context aggregation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1. Pair trace_contexts and source_span_ids per-message (via zip) instead of flattening independently. This prevents misalignment when a message has mismatched counts — orphans are dropped per-message rather than shifting all subsequent pairs out of alignment. 2. Remove TraceCapturingAggregator's override of Executor.execute() (documented as "do not override"). Capture trace data from the WorkflowContext passed to the handler instead. --- .../_workflows/_edge_runner.py | 15 ++++++--- .../packages/core/tests/workflow/test_edge.py | 31 ++++++------------- 2 files changed, 20 insertions(+), 26 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_edge_runner.py b/python/packages/core/agent_framework/_workflows/_edge_runner.py index 8955335e5b4..c14582894b9 100644 --- a/python/packages/core/agent_framework/_workflows/_edge_runner.py +++ b/python/packages/core/agent_framework/_workflows/_edge_runner.py @@ -359,13 +359,20 @@ async def send_message( # fan-in aggregation are fully preserved. Using the singular # backward-compat properties would silently drop all but the # first context per message. + # + # Pair contexts and span IDs per-message (via zip) so that a + # message with mismatched counts only drops its own orphans + # instead of shifting all subsequent pairs out of alignment + # when the flattened lists are later zipped by + # ``create_processing_span``. trace_contexts: list[dict[str, str]] = [] source_span_ids: list[str] = [] for msg in messages_to_send: - if msg.trace_contexts: - trace_contexts.extend(msg.trace_contexts) - if msg.source_span_ids: - source_span_ids.extend(msg.source_span_ids) + msg_contexts = msg.trace_contexts or [] + msg_span_ids = msg.source_span_ids or [] + for trace_context, span_id in zip(msg_contexts, msg_span_ids, strict=False): + trace_contexts.append(trace_context) + source_span_ids.append(span_id) # Create a new Message object for the aggregated data aggregated_message = WorkflowMessage( diff --git a/python/packages/core/tests/workflow/test_edge.py b/python/packages/core/tests/workflow/test_edge.py index dcb243a27c3..6d60abeb8a4 100644 --- a/python/packages/core/tests/workflow/test_edge.py +++ b/python/packages/core/tests/workflow/test_edge.py @@ -1192,7 +1192,13 @@ async def test_fan_in_edge_group_with_multiple_message_types_failed() -> None: class TraceCapturingAggregator(Executor): - """Fan-in aggregator that captures the trace contexts passed to execute().""" + """Fan-in aggregator that captures the trace contexts received by its handler. + + Captures the source trace data from the :class:`WorkflowContext` passed to + the handler rather than overriding :meth:`Executor.execute` (which is + documented as *do not override* — it owns locking, span creation, handler + dispatch, and context construction). + """ def __init__(self, *, id: str) -> None: super().__init__(id=id) @@ -1203,27 +1209,8 @@ def __init__(self, *, id: str) -> None: @handler async def mock_aggregator_handler(self, message: list[MockMessage], ctx: WorkflowContext) -> None: self.call_count += 1 - - async def execute( - self, - message: WorkflowMessage, - source_executor_ids: list[str], - state: Any, - ctx: Any, - *, - trace_contexts: list[dict[str, str]] | None = None, - source_span_ids: list[str] | None = None, - ) -> None: - self.captured_trace_contexts = trace_contexts - self.captured_source_span_ids = source_span_ids - await super().execute( - message, - source_executor_ids, - state, - ctx, - trace_contexts=trace_contexts, - source_span_ids=source_span_ids, - ) + self.captured_trace_contexts = list(ctx._trace_contexts) + self.captured_source_span_ids = list(ctx._source_span_ids) async def test_fan_in_preserves_multiple_trace_contexts_per_message() -> None: