From 71c047850c3ef0df8e7eb4ba914b1287dc335db6 Mon Sep 17 00:00:00 2001 From: Eric Zhu Date: Wed, 3 Sep 2025 21:16:01 -0700 Subject: [PATCH] Python: Add edge group traces (#597) * add edge group traces * add edge group attributes * update unit tests * add span links --- .../agent_framework_workflow/_edge_runner.py | 265 +++++++--- .../agent_framework_workflow/_telemetry.py | 96 ++++ python/packages/workflow/tests/test_edge.py | 469 ++++++++++++++++++ 3 files changed, 757 insertions(+), 73 deletions(-) diff --git a/python/packages/workflow/agent_framework_workflow/_edge_runner.py b/python/packages/workflow/agent_framework_workflow/_edge_runner.py index 7ccc50199d..ec8c6abbd6 100644 --- a/python/packages/workflow/agent_framework_workflow/_edge_runner.py +++ b/python/packages/workflow/agent_framework_workflow/_edge_runner.py @@ -10,6 +10,7 @@ from ._edge import Edge, EdgeGroup, FanInEdgeGroup, FanOutEdgeGroup, SingleEdgeG from ._executor import Executor from ._runner_context import Message, RunnerContext from ._shared_state import SharedState +from ._telemetry import EdgeGroupDeliveryStatus, workflow_tracer from ._workflow_context import WorkflowContext logger = logging.getLogger(__name__) @@ -38,7 +39,8 @@ class EdgeRunner(ABC): ctx: The context for the runner. Returns: - bool: True if the message was sent successfully, False if the target executor cannot handle the message. + bool: True if the message was processed successfully, + False if the target executor cannot handle the message. """ raise NotImplementedError @@ -49,7 +51,12 @@ class EdgeRunner(ABC): return self._executors[executor_id].can_handle(message_data) async def _execute_on_target( - self, target_id: str, source_id: str, message: Message, shared_state: SharedState, ctx: RunnerContext + self, + target_id: str, + source_ids: list[str], + message: Message, + shared_state: SharedState, + ctx: RunnerContext, ) -> None: """Execute a message on a target executor with trace context.""" if target_id not in self._executors: @@ -57,24 +64,14 @@ class EdgeRunner(ABC): target_executor = self._executors[target_id] - # Handle both old single trace context format and new multiple trace contexts format - trace_contexts = getattr(message, "trace_contexts", None) - source_span_ids = getattr(message, "source_span_ids", None) - - # Backwards compatibility: if old format is used, convert to new format - if trace_contexts is None and hasattr(message, "trace_context") and message.trace_context: - trace_contexts = [message.trace_context] - if source_span_ids is None and hasattr(message, "source_span_id") and message.source_span_id: - source_span_ids = [message.source_span_id] - # Create WorkflowContext with trace contexts from message workflow_context: WorkflowContext[Any] = WorkflowContext( target_id, - [source_id], + source_ids, shared_state, ctx, - trace_contexts=trace_contexts, # Pass trace contexts to WorkflowContext - source_span_ids=source_span_ids, # Pass source span IDs for linking + trace_contexts=message.trace_contexts, # Pass trace contexts to WorkflowContext + source_span_ids=message.source_span_ids, # Pass source span IDs for linking ) # Execute with trace context in WorkflowContext @@ -90,12 +87,47 @@ class SingleEdgeRunner(EdgeRunner): async def send_message(self, message: Message, shared_state: SharedState, ctx: RunnerContext) -> bool: """Send a message through the single edge.""" - if message.target_id and message.target_id != self._edge.target_id: - return False + should_execute = False + target_id = None + source_id = None - if self._can_handle(self._edge.target_id, message.data): - if self._edge.should_route(message.data): - await self._execute_on_target(self._edge.target_id, self._edge.source_id, message, shared_state, ctx) + with workflow_tracer.create_edge_group_processing_span( + self._edge_group.__class__.__name__, + edge_group_id=self._edge_group.id, + message_source_id=message.source_id, + message_target_id=message.target_id, + source_trace_contexts=message.trace_contexts, + source_span_ids=message.source_span_ids, + ): + try: + if message.target_id and message.target_id != self._edge.target_id: + workflow_tracer.set_edge_group_span_attributes( + False, EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH + ) + return False + + if self._can_handle(self._edge.target_id, message.data): + if self._edge.should_route(message.data): + workflow_tracer.set_edge_group_span_attributes(True, EdgeGroupDeliveryStatus.DELIVERED) + should_execute = True + target_id = self._edge.target_id + source_id = self._edge.source_id + else: + workflow_tracer.set_edge_group_span_attributes( + False, EdgeGroupDeliveryStatus.DROPPED_CONDITION_FALSE + ) + # Return True here because message was processed, just condition failed + return True + else: + workflow_tracer.set_edge_group_span_attributes(False, EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH) + return False + except Exception as e: + workflow_tracer.set_edge_group_span_attributes(False, EdgeGroupDeliveryStatus.EXCEPTION) + raise e + + # Execute outside the span + if should_execute and target_id and source_id: + await self._execute_on_target(target_id, [source_id], message, shared_state, ctx) return True return False @@ -113,37 +145,92 @@ class FanOutEdgeRunner(EdgeRunner): async def send_message(self, message: Message, shared_state: SharedState, ctx: RunnerContext) -> bool: """Send a message through all edges in the fan-out edge group.""" - selection_results = ( - self._selection_func(message.data, self._target_ids) if self._selection_func else self._target_ids - ) - if not self._validate_selection_result(selection_results): - raise RuntimeError( - f"Invalid selection result: {selection_results}. " - f"Expected selections to be a subset of valid target executor IDs: {self._target_ids}." + deliverable_edges = [] + single_target_edge = None + + # Process routing logic within span + with workflow_tracer.create_edge_group_processing_span( + self._edge_group.__class__.__name__, + edge_group_id=self._edge_group.id, + message_source_id=message.source_id, + message_target_id=message.target_id, + source_trace_contexts=message.trace_contexts, + source_span_ids=message.source_span_ids, + ): + try: + selection_results = ( + self._selection_func(message.data, self._target_ids) if self._selection_func else self._target_ids + ) + if not self._validate_selection_result(selection_results): + workflow_tracer.set_edge_group_span_attributes(False, EdgeGroupDeliveryStatus.EXCEPTION) + raise RuntimeError( + f"Invalid selection result: {selection_results}. " + f"Expected selections to be a subset of valid target executor IDs: {self._target_ids}." + ) + + if message.target_id: + # If the target ID is specified and the selection result contains it, send the message to that edge + if message.target_id in selection_results: + edge = self._target_map.get(message.target_id) + if edge and self._can_handle(edge.target_id, message.data): + if edge.should_route(message.data): + workflow_tracer.set_edge_group_span_attributes(True, EdgeGroupDeliveryStatus.DELIVERED) + single_target_edge = edge + else: + workflow_tracer.set_edge_group_span_attributes( + False, EdgeGroupDeliveryStatus.DROPPED_CONDITION_FALSE + ) + # For targeted messages with condition failure, return True (message was processed) + return True + else: + workflow_tracer.set_edge_group_span_attributes( + False, EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH + ) + # For targeted messages that can't be handled, return False + return False + else: + workflow_tracer.set_edge_group_span_attributes( + False, EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH + ) + # For targeted messages not in selection, return False + return False + else: + # If no target ID, send the message to the selected targets + for target_id in selection_results: + edge = self._target_map[target_id] + if self._can_handle(edge.target_id, message.data) and edge.should_route(message.data): + deliverable_edges.append(edge) + + if len(deliverable_edges) > 0: + workflow_tracer.set_edge_group_span_attributes(True, EdgeGroupDeliveryStatus.DELIVERED) + else: + workflow_tracer.set_edge_group_span_attributes( + False, EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH + ) + + except Exception as e: + workflow_tracer.set_edge_group_span_attributes(False, EdgeGroupDeliveryStatus.EXCEPTION) + raise e + + # Execute outside the span + if single_target_edge: + await self._execute_on_target( + single_target_edge.target_id, [single_target_edge.source_id], message, shared_state, ctx ) + return True - if message.target_id: - # If the target ID is specified and the selection result contains it, send the message to that edge - if message.target_id in selection_results: - edge = self._target_map.get(message.target_id) - if edge and self._can_handle(edge.target_id, message.data): - if edge.should_route(message.data): - await self._execute_on_target(edge.target_id, edge.source_id, message, shared_state, ctx) - return True - return False + if deliverable_edges: - # If no target ID, send the message to the selected targets - async def send_to_edge(edge: Edge) -> bool: - """Send the message to the edge.""" - if self._can_handle(edge.target_id, message.data): - if edge.should_route(message.data): - await self._execute_on_target(edge.target_id, edge.source_id, message, shared_state, ctx) + async def send_to_edge(edge: Edge) -> bool: + await self._execute_on_target(edge.target_id, [edge.source_id], message, shared_state, ctx) return True - return False - tasks = [send_to_edge(self._target_map[target_id]) for target_id in selection_results] - results = await asyncio.gather(*tasks) - return any(results) + tasks = [send_to_edge(edge) for edge in deliverable_edges] + results = await asyncio.gather(*tasks) + return any(results) + + # If we get here, it's a broadcast message with no deliverable edges + return False def _validate_selection_result(self, selection_results: list[str]) -> bool: """Validate the selection results to ensure all IDs are valid target executor IDs.""" @@ -162,40 +249,72 @@ class FanInEdgeRunner(EdgeRunner): async def send_message(self, message: Message, shared_state: SharedState, ctx: RunnerContext) -> bool: """Send a message through all edges in the fan-in edge group.""" - if message.target_id and message.target_id != self._edges[0].target_id: - return False + execution_data: dict[str, Any] | None = None - # Check if target can handle list of message data (fan-in aggregates multiple messages) - if self._can_handle(self._edges[0].target_id, [message.data]): - # If the edge can handle the data, buffer the message - self._buffer[message.source_id].append(message) - else: - # If the edge cannot handle the data, return False - return False + with workflow_tracer.create_edge_group_processing_span( + self._edge_group.__class__.__name__, + edge_group_id=self._edge_group.id, + message_source_id=message.source_id, + message_target_id=message.target_id, + source_trace_contexts=message.trace_contexts, + source_span_ids=message.source_span_ids, + ): + try: + if message.target_id and message.target_id != self._edges[0].target_id: + workflow_tracer.set_edge_group_span_attributes( + False, EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH + ) + return False - if self._is_ready_to_send(): - # If all edges in the group have data, send the buffered messages to the target executor - messages_to_send = [msg for edge in self._edges for msg in self._buffer[edge.source_id]] - self._buffer.clear() - # Send aggregated data to target - aggregated_data = [msg.data for msg in messages_to_send] + # Check if target can handle list of message data (fan-in aggregates multiple messages) + if self._can_handle(self._edges[0].target_id, [message.data]): + # If the edge can handle the data, buffer the message + self._buffer[message.source_id].append(message) + workflow_tracer.set_edge_group_span_attributes(True, EdgeGroupDeliveryStatus.BUFFERED) + else: + # If the edge cannot handle the data, return False + workflow_tracer.set_edge_group_span_attributes(False, EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH) + return False - # 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] + if self._is_ready_to_send(): + # If all edges in the group have data, prepare for execution + messages_to_send = [msg for edge in self._edges for msg in self._buffer[edge.source_id]] + self._buffer.clear() + # Send aggregated data to target + aggregated_data = [msg.data for msg in messages_to_send] - # Create a new Message object for the aggregated data - aggregated_message = Message( - data=aggregated_data, - source_id=self._edge_group.__class__.__name__, - trace_contexts=trace_contexts, - source_span_ids=source_span_ids, - ) + # 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] + + # Create a new Message object for the aggregated data + aggregated_message = Message( + data=aggregated_data, + source_id=self._edge_group.__class__.__name__, # This won't be used in self._execute_on_target. + trace_contexts=trace_contexts, + source_span_ids=source_span_ids, + ) + workflow_tracer.set_edge_group_span_attributes(True, EdgeGroupDeliveryStatus.DELIVERED) + + # Store execution data for later + execution_data = { + "target_id": self._edges[0].target_id, + "source_ids": [edge.source_id for edge in self._edges], + "message": aggregated_message, + } + + except Exception as e: + workflow_tracer.set_edge_group_span_attributes(False, EdgeGroupDeliveryStatus.EXCEPTION) + raise e + + # Execute outside the span if needed + if execution_data: await self._execute_on_target( - self._edges[0].target_id, self._edge_group.__class__.__name__, aggregated_message, shared_state, ctx + execution_data["target_id"], execution_data["source_ids"], execution_data["message"], shared_state, ctx ) + return True - return True + return True # Return True for buffered messages (waiting for more) def _is_ready_to_send(self) -> bool: """Check if all edges in the group have data to send.""" diff --git a/python/packages/workflow/agent_framework_workflow/_telemetry.py b/python/packages/workflow/agent_framework_workflow/_telemetry.py index f90d439665..625e81aa67 100644 --- a/python/packages/workflow/agent_framework_workflow/_telemetry.py +++ b/python/packages/workflow/agent_framework_workflow/_telemetry.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. +from enum import Enum from typing import TYPE_CHECKING, Any, ClassVar from agent_framework._pydantic import AFBaseSettings @@ -11,11 +12,23 @@ if TYPE_CHECKING: from ._workflow import Workflow +class EdgeGroupDeliveryStatus(Enum): + """Enum for edge group delivery status values.""" + + DELIVERED = "delivered" + DROPPED_TYPE_MISMATCH = "dropped type mismatch" + DROPPED_TARGET_MISMATCH = "dropped target mismatch" + DROPPED_CONDITION_FALSE = "dropped condition evaluated to false" + EXCEPTION = "exception" + BUFFERED = "buffered" + + # Span name constants _WORKFLOW_BUILD_SPAN = "workflow.build" _WORKFLOW_RUN_SPAN = "workflow.run" _EXECUTOR_PROCESS_SPAN = "executor.process" _MESSAGE_SEND_SPAN = "message.send" +_EDGE_GROUP_PROCESS_SPAN = "edge_group.process" class WorkflowDiagnosticSettings(AFBaseSettings): @@ -37,6 +50,7 @@ class WorkflowTracer: - Workflow execution spans (workflow.run) - Executor processing spans (executor.process) - Message sending spans (message.send) + - Edge group processing spans (edge_group.process) Implements span linking for causality without unwanted nesting. """ @@ -131,6 +145,88 @@ class WorkflowTracer: attributes=attributes, ) + def create_edge_group_processing_span( + self, + edge_group_type: str, + edge_group_id: str | None = None, + message_source_id: str | None = None, + message_target_id: str | None = None, + source_trace_contexts: list[dict[str, str]] | None = None, + source_span_ids: list[str] | None = None, + ) -> Any: + """Create an edge group processing span with optional links to source spans. + + Edge group processing spans track the processing operations in edge runners + before message delivery, including condition checking and routing decisions. + Links to source spans provide causality tracking without unwanted nesting. + + Args: + edge_group_type: The type of the edge group (class name). + edge_group_id: The unique ID of the edge group. + message_source_id: The source ID of the message being processed. + message_target_id: The target ID of the message being processed. + source_trace_contexts: Optional trace contexts from source spans for linking. + source_span_ids: Optional source span IDs for linking. + """ + attributes: dict[str, str] = { + "edge_group.type": edge_group_type, + } + + if edge_group_id is not None: + attributes["edge_group.id"] = edge_group_id + if message_source_id is not None: + attributes["message.source_id"] = message_source_id + if message_target_id is not None: + attributes["message.target_id"] = message_target_id + + # Create links to source spans for causality without nesting + links = [] + if source_trace_contexts and source_span_ids: + # Create links for all source spans (supporting fan-in with multiple sources) + for trace_context, span_id in zip(source_trace_contexts, source_span_ids, strict=False): + try: + # Extract trace and span IDs from the trace context + # This is a simplified approach - in production you'd want more robust parsing + traceparent = trace_context.get("traceparent", "") + if traceparent: + # traceparent format: "00-{trace_id}-{parent_span_id}-{trace_flags}" + parts = traceparent.split("-") + if len(parts) >= 3: + trace_id_hex = parts[1] + # Use the source_span_id that was saved from the publishing span + + # Create span context for linking + span_context = SpanContext( + trace_id=int(trace_id_hex, 16), + span_id=int(span_id, 16), + is_remote=True, + ) + links.append(Link(span_context)) + except (ValueError, TypeError, AttributeError): + # If linking fails, continue without link (graceful degradation) + pass + + return self.tracer.start_as_current_span( + _EDGE_GROUP_PROCESS_SPAN, + kind=SpanKind.INTERNAL, + attributes=attributes, + links=links, + ) + + def set_edge_group_span_attributes(self, delivered: bool, delivery_status: EdgeGroupDeliveryStatus) -> None: + """Set edge group span attributes for delivery status. + + Args: + delivered: Whether the message was delivered. + delivery_status: The delivery status from EdgeGroupDeliveryStatus enum. + """ + span = get_current_span() + if span and span.is_recording(): + span.set_attributes({ + "edge_group.delivered": delivered, + "edge_group.delivery_status": delivery_status.value, + }) + def add_workflow_event(self, event_name: str, attributes: Attributes | None = None) -> None: """Add an event to the current workflow span. diff --git a/python/packages/workflow/tests/test_edge.py b/python/packages/workflow/tests/test_edge.py index 3e452e95ab..98ecee5138 100644 --- a/python/packages/workflow/tests/test_edge.py +++ b/python/packages/workflow/tests/test_edge.py @@ -6,6 +6,9 @@ from unittest.mock import patch import pytest from agent_framework.workflow import Executor, WorkflowContext, handler +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from agent_framework_workflow._edge import ( Edge, @@ -19,6 +22,54 @@ from agent_framework_workflow._edge import ( from agent_framework_workflow._edge_runner import create_edge_runner from agent_framework_workflow._runner_context import InProcRunnerContext, Message from agent_framework_workflow._shared_state import SharedState +from agent_framework_workflow._telemetry import EdgeGroupDeliveryStatus, workflow_tracer + + +@pytest.fixture +def tracing_enabled(): + """Enable tracing for tests.""" + import os + + original_value = os.environ.get("AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS") + os.environ["AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS"] = "true" + + # Force reload the settings to pick up the environment variable + from agent_framework_workflow._telemetry import WorkflowDiagnosticSettings + + workflow_tracer.settings = WorkflowDiagnosticSettings() + + yield + + # Restore original value + if original_value is None: + os.environ.pop("AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS", None) + else: + os.environ["AGENT_FRAMEWORK_WORKFLOW_ENABLE_OTEL_DIAGNOSTICS"] = original_value + + # Reload settings again + workflow_tracer.settings = WorkflowDiagnosticSettings() + + +@pytest.fixture +def span_exporter(tracing_enabled): + """Set up OpenTelemetry test infrastructure.""" + + # Use the built-in InMemorySpanExporter for better compatibility + exporter = InMemorySpanExporter() + tracer_provider = TracerProvider() + tracer_provider.add_span_processor(SimpleSpanProcessor(exporter)) + + # Store original tracer + original_tracer = workflow_tracer.tracer + + # Set up our test tracer + workflow_tracer.tracer = tracer_provider.get_tracer("agent_framework") + + yield exporter + + # Clean up + exporter.clear() + workflow_tracer.tracer = original_tracer @dataclass @@ -207,6 +258,200 @@ async def test_single_edge_group_send_message_with_invalid_data() -> None: assert success is False +async def test_single_edge_group_send_message_with_condition_pass() -> None: + """Test sending a message through a single edge runner with a condition that passes.""" + source = MockExecutor(id="source_executor") + target = MockExecutor(id="target_executor") + + executors: dict[str, Executor] = {source.id: source, target.id: target} + # Create edge group with condition that passes when data == "test" + edge_group = SingleEdgeGroup(source_id=source.id, target_id=target.id, condition=lambda x: x.data == "test") + + edge_runner = create_edge_runner(edge_group, executors) + shared_state = SharedState() + ctx = InProcRunnerContext() + + data = MockMessage(data="test") + message = Message(data=data, source_id=source.id) + + success = await edge_runner.send_message(message, shared_state, ctx) + assert success is True + assert target.call_count == 1 + assert target.last_message.data == "test" + + +async def test_single_edge_group_send_message_with_condition_fail() -> None: + """Test sending a message through a single edge runner with a condition that fails.""" + source = MockExecutor(id="source_executor") + target = MockExecutor(id="target_executor") + + executors: dict[str, Executor] = {source.id: source, target.id: target} + # Create edge group with condition that passes when data == "test" + edge_group = SingleEdgeGroup(source_id=source.id, target_id=target.id, condition=lambda x: x.data == "test") + + edge_runner = create_edge_runner(edge_group, executors) + shared_state = SharedState() + ctx = InProcRunnerContext() + + data = MockMessage(data="different") + message = Message(data=data, source_id=source.id) + + success = await edge_runner.send_message(message, shared_state, ctx) + # Should return True because message was processed, but condition failed + assert success is True + # Target should not be called because condition failed + assert target.call_count == 0 + + +async def test_single_edge_group_tracing_success(span_exporter) -> None: + """Test that single edge group processing creates proper success spans.""" + source = MockExecutor(id="source_executor") + target = MockExecutor(id="target_executor") + + executors: dict[str, Executor] = {source.id: source, target.id: target} + edge_group = SingleEdgeGroup(source_id=source.id, target_id=target.id) + + edge_runner = create_edge_runner(edge_group, executors) + shared_state = SharedState() + ctx = InProcRunnerContext() + + # Create trace context and span IDs to simulate a message with tracing information + trace_contexts = [{"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"}] + source_span_ids = ["00f067aa0ba902b7"] + + data = MockMessage(data="test") + message = Message(data=data, source_id=source.id, trace_contexts=trace_contexts, source_span_ids=source_span_ids) + + # Clear any build spans + span_exporter.clear() + + success = await edge_runner.send_message(message, shared_state, ctx) + assert success is True + + spans = span_exporter.get_finished_spans() + edge_group_spans = [s for s in spans if s.name == "edge_group.process"] + + assert len(edge_group_spans) == 1 + + span = edge_group_spans[0] + assert span.attributes is not None + assert span.attributes.get("edge_group.type") == "SingleEdgeGroup" + assert span.attributes.get("edge_group.delivered") is True + assert span.attributes.get("edge_group.delivery_status") == EdgeGroupDeliveryStatus.DELIVERED.value + assert span.attributes.get("edge_group.id") is not None + assert span.attributes.get("message.source_id") == source.id + + # Verify span links are created + assert span.links is not None + assert len(span.links) == 1 + + link = span.links[0] + # Verify the link points to the correct trace and span + assert link.context.trace_id == int("4bf92f3577b34da6a3ce929d0e0e4736", 16) + assert link.context.span_id == int("00f067aa0ba902b7", 16) + + +async def test_single_edge_group_tracing_condition_failure(span_exporter) -> None: + """Test that single edge group processing creates proper spans for condition failures.""" + source = MockExecutor(id="source_executor") + target = MockExecutor(id="target_executor") + + executors: dict[str, Executor] = {source.id: source, target.id: target} + edge_group = SingleEdgeGroup(source_id=source.id, target_id=target.id, condition=lambda x: x.data == "pass") + + edge_runner = create_edge_runner(edge_group, executors) + shared_state = SharedState() + ctx = InProcRunnerContext() + + data = MockMessage(data="fail") + message = Message(data=data, source_id=source.id) + + # Clear any build spans + span_exporter.clear() + + success = await edge_runner.send_message(message, shared_state, ctx) + assert success is True # Returns True but condition failed + + spans = span_exporter.get_finished_spans() + edge_group_spans = [s for s in spans if s.name == "edge_group.process"] + + assert len(edge_group_spans) == 1 + + span = edge_group_spans[0] + assert span.attributes is not None + assert span.attributes.get("edge_group.type") == "SingleEdgeGroup" + assert span.attributes.get("edge_group.delivered") is False + assert span.attributes.get("edge_group.delivery_status") == EdgeGroupDeliveryStatus.DROPPED_CONDITION_FALSE.value + + +async def test_single_edge_group_tracing_type_mismatch(span_exporter) -> None: + """Test that single edge group processing creates proper spans for type mismatches.""" + source = MockExecutor(id="source_executor") + target = MockExecutor(id="target_executor") + + executors: dict[str, Executor] = {source.id: source, target.id: target} + edge_group = SingleEdgeGroup(source_id=source.id, target_id=target.id) + + edge_runner = create_edge_runner(edge_group, executors) + shared_state = SharedState() + ctx = InProcRunnerContext() + + # Send incompatible data type + data = "invalid_data" + message = Message(data=data, source_id=source.id) + + # Clear any build spans + span_exporter.clear() + + success = await edge_runner.send_message(message, shared_state, ctx) + assert success is False + + spans = span_exporter.get_finished_spans() + edge_group_spans = [s for s in spans if s.name == "edge_group.process"] + + assert len(edge_group_spans) == 1 + + span = edge_group_spans[0] + assert span.attributes is not None + assert span.attributes.get("edge_group.type") == "SingleEdgeGroup" + assert span.attributes.get("edge_group.delivered") is False + assert span.attributes.get("edge_group.delivery_status") == EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value + + +async def test_single_edge_group_tracing_target_mismatch(span_exporter) -> None: + """Test that single edge group processing creates proper spans for target mismatches.""" + source = MockExecutor(id="source_executor") + target = MockExecutor(id="target_executor") + + executors: dict[str, Executor] = {source.id: source, target.id: target} + edge_group = SingleEdgeGroup(source_id=source.id, target_id=target.id) + + edge_runner = create_edge_runner(edge_group, executors) + shared_state = SharedState() + ctx = InProcRunnerContext() + + data = MockMessage(data="test") + message = Message(data=data, source_id=source.id, target_id="wrong_target") + + # Clear any build spans + span_exporter.clear() + + success = await edge_runner.send_message(message, shared_state, ctx) + assert success is False + + spans = span_exporter.get_finished_spans() + edge_group_spans = [s for s in spans if s.name == "edge_group.process"] + + assert len(edge_group_spans) == 1 + + span = edge_group_spans[0] + assert span.attributes is not None + assert span.attributes.get("edge_group.type") == "SingleEdgeGroup" + assert span.attributes.get("edge_group.delivered") is False + assert span.attributes.get("edge_group.delivery_status") == EdgeGroupDeliveryStatus.DROPPED_TARGET_MISMATCH.value + assert span.attributes.get("message.target_id") == "wrong_target" + + # endregion SingleEdgeGroup @@ -526,6 +771,109 @@ async def test_source_edge_group_with_selection_func_send_message_with_target_in assert success is False +async def test_fan_out_edge_group_tracing_success(span_exporter) -> None: + """Test that fan-out edge group processing creates proper success spans.""" + source = MockExecutor(id="source_executor") + target1 = MockExecutor(id="target_executor_1") + target2 = MockExecutor(id="target_executor_2") + + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_group = FanOutEdgeGroup(source_id=source.id, target_ids=[target1.id, target2.id]) + + edge_runner = create_edge_runner(edge_group, executors) + shared_state = SharedState() + ctx = InProcRunnerContext() + + # Create trace context and span IDs to simulate a message with tracing information + trace_contexts = [{"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"}] + source_span_ids = ["00f067aa0ba902b7"] + + data = MockMessage(data="test") + message = Message(data=data, source_id=source.id, trace_contexts=trace_contexts, source_span_ids=source_span_ids) + + # Clear any build spans + span_exporter.clear() + + success = await edge_runner.send_message(message, shared_state, ctx) + assert success is True + + spans = span_exporter.get_finished_spans() + edge_group_spans = [s for s in spans if s.name == "edge_group.process"] + + assert len(edge_group_spans) == 1 + + span = edge_group_spans[0] + assert span.attributes is not None + assert span.attributes.get("edge_group.type") == "FanOutEdgeGroup" + assert span.attributes.get("edge_group.delivered") is True + assert span.attributes.get("edge_group.delivery_status") == EdgeGroupDeliveryStatus.DELIVERED.value + assert span.attributes.get("edge_group.id") is not None + assert span.attributes.get("message.source_id") == source.id + + # Verify span links are created + assert span.links is not None + assert len(span.links) == 1 + + link = span.links[0] + # Verify the link points to the correct trace and span + assert link.context.trace_id == int("4bf92f3577b34da6a3ce929d0e0e4736", 16) + assert link.context.span_id == int("00f067aa0ba902b7", 16) + + +async def test_fan_out_edge_group_tracing_with_target(span_exporter) -> None: + """Test that fan-out edge group processing creates proper spans for targeted messages.""" + source = MockExecutor(id="source_executor") + target1 = MockExecutor(id="target_executor_1") + target2 = MockExecutor(id="target_executor_2") + + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_group = FanOutEdgeGroup(source_id=source.id, target_ids=[target1.id, target2.id]) + + edge_runner = create_edge_runner(edge_group, executors) + shared_state = SharedState() + ctx = InProcRunnerContext() + + # Create trace context and span IDs to simulate a message with tracing information + trace_contexts = [{"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"}] + source_span_ids = ["00f067aa0ba902b7"] + + data = MockMessage(data="test") + message = Message( + data=data, + source_id=source.id, + target_id=target1.id, + trace_contexts=trace_contexts, + source_span_ids=source_span_ids, + ) + + # Clear any build spans + span_exporter.clear() + + success = await edge_runner.send_message(message, shared_state, ctx) + assert success is True + + spans = span_exporter.get_finished_spans() + edge_group_spans = [s for s in spans if s.name == "edge_group.process"] + + assert len(edge_group_spans) == 1 + + span = edge_group_spans[0] + assert span.attributes is not None + assert span.attributes.get("edge_group.type") == "FanOutEdgeGroup" + assert span.attributes.get("edge_group.delivered") is True + assert span.attributes.get("edge_group.delivery_status") == EdgeGroupDeliveryStatus.DELIVERED.value + assert span.attributes.get("message.target_id") == target1.id + + # Verify span links are created + assert span.links is not None + assert len(span.links) == 1 + + link = span.links[0] + # Verify the link points to the correct trace and span + assert link.context.trace_id == int("4bf92f3577b34da6a3ce929d0e0e4736", 16) + assert link.context.span_id == int("00f067aa0ba902b7", 16) + + # endregion FanOutEdgeGroup # region FanInEdgeGroup @@ -638,6 +986,127 @@ async def test_target_edge_group_send_message_with_invalid_data() -> None: assert success is False +async def test_fan_in_edge_group_tracing_buffered(span_exporter) -> None: + """Test that fan-in edge group processing creates proper spans for buffered messages.""" + source1 = MockExecutor(id="source_executor_1") + source2 = MockExecutor(id="source_executor_2") + target = MockAggregator(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) + shared_state = SharedState() + ctx = InProcRunnerContext() + + data = MockMessage(data="test") + + # Create trace context and span IDs to simulate a message with tracing information + trace_contexts1 = [{"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"}] + source_span_ids1 = ["00f067aa0ba902b7"] + + trace_contexts2 = [{"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b8-01"}] + source_span_ids2 = ["00f067aa0ba902b8"] + + # Clear any build spans + span_exporter.clear() + + # Send first message (should be buffered) + success = await edge_runner.send_message( + Message(data=data, source_id=source1.id, trace_contexts=trace_contexts1, source_span_ids=source_span_ids1), + shared_state, + ctx, + ) + assert success is True + + spans = span_exporter.get_finished_spans() + edge_group_spans = [s for s in spans if s.name == "edge_group.process"] + + assert len(edge_group_spans) == 1 + + span = edge_group_spans[0] + assert span.attributes is not None + assert span.attributes.get("edge_group.type") == "FanInEdgeGroup" + assert span.attributes.get("edge_group.delivered") is True + assert span.attributes.get("edge_group.delivery_status") == EdgeGroupDeliveryStatus.BUFFERED.value + assert span.attributes.get("message.source_id") == source1.id + + # Verify span links are created for first message + assert span.links is not None + assert len(span.links) == 1 + + link = span.links[0] + # Verify the link points to the correct trace and span + assert link.context.trace_id == int("4bf92f3577b34da6a3ce929d0e0e4736", 16) + assert link.context.span_id == int("00f067aa0ba902b7", 16) + + # Clear spans and send second message (should trigger delivery) + span_exporter.clear() + + success = await edge_runner.send_message( + Message(data=data, source_id=source2.id, trace_contexts=trace_contexts2, source_span_ids=source_span_ids2), + shared_state, + ctx, + ) + assert success is True + + spans = span_exporter.get_finished_spans() + edge_group_spans = [s for s in spans if s.name == "edge_group.process"] + + assert len(edge_group_spans) == 1 + + span = edge_group_spans[0] + assert span.attributes is not None + assert span.attributes.get("edge_group.type") == "FanInEdgeGroup" + assert span.attributes.get("edge_group.delivered") is True + assert span.attributes.get("edge_group.delivery_status") == EdgeGroupDeliveryStatus.DELIVERED.value + assert span.attributes.get("message.source_id") == source2.id + + # Verify span links are created for second message + assert span.links is not None + assert len(span.links) == 1 + + link = span.links[0] + # Verify the link points to the correct trace and span for the second message + assert link.context.trace_id == int("4bf92f3577b34da6a3ce929d0e0e4736", 16) + assert link.context.span_id == int("00f067aa0ba902b8", 16) + + +async def test_fan_in_edge_group_tracing_type_mismatch(span_exporter) -> None: + """Test that fan-in edge group processing creates proper spans for type mismatches.""" + source1 = MockExecutor(id="source_executor_1") + source2 = MockExecutor(id="source_executor_2") + target = MockAggregator(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) + shared_state = SharedState() + ctx = InProcRunnerContext() + + # Send incompatible data type + data = "invalid_data" + message = Message(data=data, source_id=source1.id) + + # Clear any build spans + span_exporter.clear() + + success = await edge_runner.send_message(message, shared_state, ctx) + assert success is False + + spans = span_exporter.get_finished_spans() + edge_group_spans = [s for s in spans if s.name == "edge_group.process"] + + assert len(edge_group_spans) == 1 + + span = edge_group_spans[0] + assert span.attributes is not None + assert span.attributes.get("edge_group.type") == "FanInEdgeGroup" + assert span.attributes.get("edge_group.delivered") is False + assert span.attributes.get("edge_group.delivery_status") == EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value + + # endregion FanInEdgeGroup # region SwitchCaseEdgeGroup