mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: Tracing for workflows (#480)
* workflow tracing design doc * add tracing implementation for workflow * fix bug caused by double wrapping of sub workflow request * add unit tests for tracing * add documentation for workflow tracing * remove unnecessary file * update aspire command * fix tests * proper serialization of subworkflows and add workflow.definition * add serialization test * fix subworkflow serialization * workflow_id --> id * update workflow sample to address comments * update naming; use costant * use NoOpTracer instead of nullcontext * use span event instead of attribtutes for status * fix typing * add workflow.build span * rename methods for clarity * ensure all source trace contexts are propagated in fan in
This commit is contained in:
Eric Zhu
committed by
GitHub
Unverified
GitHub
parent
78125f019a
commit
379e3b9a00
@@ -15,6 +15,9 @@ from agent_framework_workflow._edge import (
|
||||
SwitchCaseEdgeGroupCase,
|
||||
SwitchCaseEdgeGroupDefault,
|
||||
)
|
||||
from agent_framework_workflow._executor import (
|
||||
WorkflowExecutor,
|
||||
)
|
||||
|
||||
|
||||
class SampleExecutor(Executor):
|
||||
@@ -398,6 +401,110 @@ class TestSerializationWorkflowClasses:
|
||||
"JSON SwitchCaseEdgeGroup edges should not have condition_name"
|
||||
)
|
||||
|
||||
def test_nested_workflow_executor_serialization(self) -> None:
|
||||
"""Test complete serialization of deeply nested WorkflowExecutors (subworkflows within subworkflows).
|
||||
|
||||
This test verifies that nested WorkflowExecutor objects are fully serialized with their
|
||||
complete workflow structures, including deeply nested workflows and all their executors.
|
||||
"""
|
||||
# Create innermost workflow
|
||||
inner_executor = SampleExecutor(id="inner-exec")
|
||||
inner_workflow = WorkflowBuilder().set_start_executor(inner_executor).set_max_iterations(10).build()
|
||||
|
||||
# Create middle workflow with WorkflowExecutor
|
||||
inner_workflow_executor = WorkflowExecutor(workflow=inner_workflow, id="inner-workflow-exec")
|
||||
middle_executor = SampleExecutor(id="middle-exec")
|
||||
middle_workflow = (
|
||||
WorkflowBuilder()
|
||||
.set_start_executor(middle_executor)
|
||||
.add_edge(middle_executor, inner_workflow_executor)
|
||||
.set_max_iterations(20)
|
||||
.build()
|
||||
)
|
||||
|
||||
# Create outer workflow with nested WorkflowExecutor
|
||||
middle_workflow_executor = WorkflowExecutor(workflow=middle_workflow, id="middle-workflow-exec")
|
||||
outer_executor = SampleExecutor(id="outer-exec")
|
||||
outer_workflow = (
|
||||
WorkflowBuilder()
|
||||
.set_start_executor(outer_executor)
|
||||
.add_edge(outer_executor, middle_workflow_executor)
|
||||
.set_max_iterations(30)
|
||||
.build()
|
||||
)
|
||||
|
||||
# Test serialization of the nested structure
|
||||
data = outer_workflow.model_dump()
|
||||
|
||||
# Verify outer structure
|
||||
assert data["start_executor_id"] == "outer-exec"
|
||||
assert data["max_iterations"] == 30
|
||||
assert "outer-exec" in data["executors"]
|
||||
assert "middle-workflow-exec" in data["executors"]
|
||||
|
||||
# Verify middle WorkflowExecutor is present with full nested workflow serialization
|
||||
middle_exec_data = data["executors"]["middle-workflow-exec"]
|
||||
assert middle_exec_data["type"] == "WorkflowExecutor"
|
||||
assert middle_exec_data["id"] == "middle-workflow-exec"
|
||||
|
||||
# Verify the nested workflow is fully serialized
|
||||
assert "workflow" in middle_exec_data, "WorkflowExecutor should include nested workflow in serialization"
|
||||
middle_workflow_data = middle_exec_data["workflow"]
|
||||
assert "start_executor_id" in middle_workflow_data
|
||||
assert "executors" in middle_workflow_data
|
||||
assert "max_iterations" in middle_workflow_data
|
||||
assert middle_workflow_data["start_executor_id"] == "middle-exec"
|
||||
assert middle_workflow_data["max_iterations"] == 20
|
||||
|
||||
# Verify the deeply nested executors are present
|
||||
assert "middle-exec" in middle_workflow_data["executors"]
|
||||
assert "inner-workflow-exec" in middle_workflow_data["executors"]
|
||||
|
||||
# Verify the innermost WorkflowExecutor is also fully serialized
|
||||
inner_workflow_exec_data = middle_workflow_data["executors"]["inner-workflow-exec"]
|
||||
assert inner_workflow_exec_data["type"] == "WorkflowExecutor"
|
||||
assert "workflow" in inner_workflow_exec_data, "Deeply nested WorkflowExecutor should also include its workflow"
|
||||
innermost_workflow_data = inner_workflow_exec_data["workflow"]
|
||||
assert "start_executor_id" in innermost_workflow_data
|
||||
assert "executors" in innermost_workflow_data
|
||||
assert "max_iterations" in innermost_workflow_data
|
||||
assert innermost_workflow_data["start_executor_id"] == "inner-exec"
|
||||
assert innermost_workflow_data["max_iterations"] == 10
|
||||
assert "inner-exec" in innermost_workflow_data["executors"]
|
||||
|
||||
# Test JSON serialization preserves the complete nested structure
|
||||
json_str = outer_workflow.model_dump_json()
|
||||
parsed = json.loads(json_str)
|
||||
|
||||
# Verify the complete structure is preserved in JSON
|
||||
middle_exec_json = parsed["executors"]["middle-workflow-exec"]
|
||||
assert middle_exec_json["type"] == "WorkflowExecutor"
|
||||
assert middle_exec_json["id"] == "middle-workflow-exec"
|
||||
|
||||
# Verify nested workflow is present in JSON
|
||||
assert "workflow" in middle_exec_json, "JSON serialization should include nested workflow"
|
||||
middle_workflow_json = middle_exec_json["workflow"]
|
||||
assert middle_workflow_json["start_executor_id"] == "middle-exec"
|
||||
assert middle_workflow_json["max_iterations"] == 20
|
||||
assert "middle-exec" in middle_workflow_json["executors"]
|
||||
assert "inner-workflow-exec" in middle_workflow_json["executors"]
|
||||
|
||||
# Verify deeply nested structure in JSON
|
||||
inner_workflow_exec_json = middle_workflow_json["executors"]["inner-workflow-exec"]
|
||||
assert inner_workflow_exec_json["type"] == "WorkflowExecutor"
|
||||
assert "workflow" in inner_workflow_exec_json, "Deeply nested WorkflowExecutor should be in JSON"
|
||||
innermost_workflow_json = inner_workflow_exec_json["workflow"]
|
||||
assert innermost_workflow_json["start_executor_id"] == "inner-exec"
|
||||
assert innermost_workflow_json["max_iterations"] == 10
|
||||
assert "inner-exec" in innermost_workflow_json["executors"]
|
||||
|
||||
# Test that WorkflowExecutor also serializes correctly when accessed directly
|
||||
direct_middle_data = middle_workflow_executor.model_dump()
|
||||
assert "workflow" in direct_middle_data
|
||||
assert direct_middle_data["type"] == "WorkflowExecutor"
|
||||
assert "executors" in direct_middle_data["workflow"]
|
||||
assert "inner-workflow-exec" in direct_middle_data["workflow"]["executors"]
|
||||
|
||||
def test_switch_case_edge_group_serialization_with_named_condition(self) -> None:
|
||||
"""Test that SwitchCaseEdgeGroup with named condition function serializes condition_name correctly."""
|
||||
|
||||
@@ -440,7 +547,7 @@ class TestSerializationWorkflowClasses:
|
||||
assert "executors" in data
|
||||
assert "start_executor_id" in data
|
||||
assert "max_iterations" in data
|
||||
assert "workflow_id" in data
|
||||
assert "id" in data
|
||||
|
||||
assert data["start_executor_id"] == "executor1"
|
||||
assert "executor1" in data["executors"]
|
||||
|
||||
@@ -0,0 +1,532 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import os
|
||||
from collections.abc import Generator
|
||||
from typing import Any, cast
|
||||
|
||||
import pytest
|
||||
from opentelemetry import trace
|
||||
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 import WorkflowBuilder
|
||||
from agent_framework_workflow._executor import Executor, handler
|
||||
from agent_framework_workflow._runner_context import InProcRunnerContext, Message
|
||||
from agent_framework_workflow._shared_state import SharedState
|
||||
from agent_framework_workflow._telemetry import WorkflowTracer, workflow_tracer
|
||||
from agent_framework_workflow._workflow import Workflow
|
||||
from agent_framework_workflow._workflow_context import WorkflowContext
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tracing_enabled() -> Generator[None, None, None]:
|
||||
"""Enable tracing for tests."""
|
||||
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: Any) -> Generator[InMemorySpanExporter, None, None]:
|
||||
"""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
|
||||
|
||||
|
||||
class MockExecutor(Executor):
|
||||
"""Mock executor for testing."""
|
||||
|
||||
def __init__(self, id: str = "mock_executor") -> None:
|
||||
super().__init__(id=id)
|
||||
# Use private field to avoid Pydantic validation
|
||||
self._processed_messages: list[str] = []
|
||||
|
||||
@handler
|
||||
async def handle_message(self, message: str, ctx: WorkflowContext[str]) -> None:
|
||||
"""Handle string messages."""
|
||||
self._processed_messages.append(message)
|
||||
await ctx.send_message(f"processed: {message}")
|
||||
|
||||
@property
|
||||
def processed_messages(self) -> list[str]:
|
||||
"""Access to processed messages for testing."""
|
||||
return self._processed_messages
|
||||
|
||||
|
||||
class SecondExecutor(Executor):
|
||||
"""Second executor for testing message chains."""
|
||||
|
||||
def __init__(self, id: str = "second_executor") -> None:
|
||||
super().__init__(id=id)
|
||||
# Use private field to avoid Pydantic validation
|
||||
self._processed_messages: list[str] = []
|
||||
|
||||
@handler
|
||||
async def handle_message(self, message: str, ctx: WorkflowContext[None]) -> None:
|
||||
"""Handle string messages."""
|
||||
self._processed_messages.append(message)
|
||||
|
||||
@property
|
||||
def processed_messages(self) -> list[str]:
|
||||
"""Access to processed messages for testing."""
|
||||
return self._processed_messages
|
||||
|
||||
|
||||
class ProcessingExecutor(Executor):
|
||||
"""Executor that processes and forwards messages with a custom prefix."""
|
||||
|
||||
def __init__(self, id: str, prefix: str = "processed") -> None:
|
||||
super().__init__(id=id)
|
||||
# Use private field to avoid Pydantic validation
|
||||
self._processed_messages: list[str] = []
|
||||
self._prefix = prefix
|
||||
|
||||
@handler
|
||||
async def handle_message(self, message: str, ctx: WorkflowContext[str]) -> None:
|
||||
"""Handle string messages and send them forward with prefix."""
|
||||
self._processed_messages.append(message)
|
||||
await ctx.send_message(f"{self._prefix}: {message}")
|
||||
|
||||
@property
|
||||
def processed_messages(self) -> list[str]:
|
||||
return self._processed_messages
|
||||
|
||||
|
||||
class FanInAggregator(Executor):
|
||||
"""Fan-in aggregator that expects a list of inputs."""
|
||||
|
||||
def __init__(self, id: str = "aggregator") -> None:
|
||||
super().__init__(id=id)
|
||||
# Use private field to avoid Pydantic validation
|
||||
self._processed_messages: list[Any] = []
|
||||
|
||||
@handler
|
||||
async def handle_aggregated_data(self, messages: list[str], ctx: WorkflowContext[None]) -> None:
|
||||
# Process aggregated messages from fan-in
|
||||
aggregated = f"aggregated: {', '.join(messages)}"
|
||||
self._processed_messages.append(aggregated)
|
||||
|
||||
@property
|
||||
def processed_messages(self) -> list[Any]:
|
||||
"""Access to processed messages for testing."""
|
||||
return self._processed_messages
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_workflow_tracer_configuration() -> None:
|
||||
"""Test that workflow tracer can be enabled and disabled."""
|
||||
# Test disabled by default
|
||||
tracer = WorkflowTracer()
|
||||
assert not tracer.enabled
|
||||
|
||||
# Test enabled with environment variable
|
||||
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
|
||||
|
||||
tracer.settings = WorkflowDiagnosticSettings()
|
||||
|
||||
assert tracer.enabled
|
||||
|
||||
# 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
|
||||
tracer.settings = WorkflowDiagnosticSettings()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_span_creation_and_attributes(tracing_enabled: Any, span_exporter: InMemorySpanExporter) -> None:
|
||||
"""Test creation and attributes of all span types (workflow, processing, sending)."""
|
||||
# Create a mock workflow object
|
||||
mock_workflow = cast(
|
||||
Workflow,
|
||||
type(
|
||||
"MockWorkflow",
|
||||
(),
|
||||
{
|
||||
"id": "test-workflow-123",
|
||||
"max_iterations": 100,
|
||||
"model_dump_json": lambda self: '{"id": "test-workflow-123", "type": "mock"}',
|
||||
},
|
||||
)(),
|
||||
)
|
||||
|
||||
# Test all span types in nested context
|
||||
with workflow_tracer.create_workflow_run_span(mock_workflow) as workflow_span:
|
||||
workflow_tracer.add_workflow_event("workflow.started")
|
||||
|
||||
with (
|
||||
workflow_tracer.create_processing_span("executor-456", "TestExecutor", "TestMessage") as processing_span,
|
||||
workflow_tracer.create_sending_span("ResponseMessage", "target-789") as sending_span,
|
||||
):
|
||||
# Verify all spans are recording
|
||||
assert workflow_span is not None and workflow_span.is_recording()
|
||||
assert processing_span is not None and processing_span.is_recording()
|
||||
assert sending_span is not None and sending_span.is_recording()
|
||||
|
||||
spans = span_exporter.get_finished_spans()
|
||||
assert len(spans) == 3
|
||||
|
||||
# Check workflow span
|
||||
workflow_span = next(s for s in spans if s.name == "workflow.run")
|
||||
assert workflow_span.kind == trace.SpanKind.INTERNAL
|
||||
assert workflow_span.attributes is not None
|
||||
assert workflow_span.attributes.get("workflow.id") == "test-workflow-123"
|
||||
assert workflow_span.events is not None
|
||||
event_names = [event.name for event in workflow_span.events]
|
||||
assert "workflow.started" in event_names
|
||||
|
||||
# Check processing span
|
||||
processing_span = next(s for s in spans if s.name == "executor.process")
|
||||
assert processing_span.kind == trace.SpanKind.INTERNAL
|
||||
assert processing_span.attributes is not None
|
||||
assert processing_span.attributes.get("executor.id") == "executor-456"
|
||||
assert processing_span.attributes.get("executor.type") == "TestExecutor"
|
||||
assert processing_span.attributes.get("message.type") == "TestMessage"
|
||||
|
||||
# Check sending span
|
||||
sending_span = next(s for s in spans if s.name == "message.send")
|
||||
assert sending_span.kind == trace.SpanKind.PRODUCER
|
||||
assert sending_span.attributes is not None
|
||||
assert sending_span.attributes.get("message.type") == "ResponseMessage"
|
||||
assert sending_span.attributes.get("message.destination_executor_id") == "target-789"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trace_context_handling(tracing_enabled: Any, span_exporter: InMemorySpanExporter) -> None:
|
||||
"""Test trace context propagation and handling in messages and executors."""
|
||||
shared_state = SharedState()
|
||||
ctx = InProcRunnerContext()
|
||||
executor = MockExecutor("test-executor")
|
||||
|
||||
# Test trace context propagation in messages
|
||||
workflow_ctx: WorkflowContext[str] = WorkflowContext(
|
||||
"test-executor",
|
||||
["source"],
|
||||
shared_state,
|
||||
ctx,
|
||||
trace_contexts=[{"traceparent": "00-12345678901234567890123456789012-1234567890123456-01"}],
|
||||
source_span_ids=["1234567890123456"],
|
||||
)
|
||||
|
||||
# Send a message (this should create a sending span and propagate trace context)
|
||||
await workflow_ctx.send_message("test message")
|
||||
|
||||
# Check that message was created with trace context
|
||||
messages = await ctx.drain_messages()
|
||||
assert len(messages) == 1
|
||||
message_list = list(messages.values())[0]
|
||||
assert len(message_list) == 1
|
||||
message = message_list[0]
|
||||
assert message.trace_context is not None
|
||||
assert message.source_span_id is not None
|
||||
|
||||
# Test executor trace context handling
|
||||
await executor.execute("test message", workflow_ctx)
|
||||
|
||||
# Check that spans were created with proper attributes
|
||||
spans = span_exporter.get_finished_spans()
|
||||
processing_spans = [s for s in spans if s.name == "executor.process"]
|
||||
sending_spans = [s for s in spans if s.name == "message.send"]
|
||||
|
||||
assert len(processing_spans) >= 1
|
||||
assert len(sending_spans) >= 1
|
||||
|
||||
# Verify processing span attributes
|
||||
processing_span = processing_spans[0]
|
||||
assert processing_span.attributes is not None
|
||||
assert processing_span.attributes.get("executor.id") == "test-executor"
|
||||
assert processing_span.attributes.get("executor.type") == "MockExecutor"
|
||||
assert processing_span.attributes.get("message.type") == "str"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trace_context_disabled_when_tracing_disabled() -> None:
|
||||
"""Test that no trace context is added when tracing is disabled."""
|
||||
# Tracing should be disabled by default
|
||||
shared_state = SharedState()
|
||||
ctx = InProcRunnerContext()
|
||||
|
||||
workflow_ctx: WorkflowContext[str] = WorkflowContext(
|
||||
"test-executor",
|
||||
["source"],
|
||||
shared_state,
|
||||
ctx,
|
||||
)
|
||||
|
||||
# Send a message
|
||||
await workflow_ctx.send_message("test message")
|
||||
|
||||
# Check that message was created without trace context
|
||||
messages = await ctx.drain_messages()
|
||||
message = list(messages.values())[0][0]
|
||||
|
||||
# When tracing is disabled, trace_context should be None
|
||||
assert message.trace_context is None
|
||||
assert message.source_span_id is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_end_to_end_workflow_tracing(tracing_enabled: Any, span_exporter: InMemorySpanExporter) -> None:
|
||||
"""Test end-to-end tracing including workflow build, execution, and span linking with fan-in edges."""
|
||||
# Create executors for fan-in scenario
|
||||
executor1 = MockExecutor("executor1")
|
||||
executor2 = ProcessingExecutor("executor2", "second")
|
||||
executor3 = ProcessingExecutor("executor3", "third")
|
||||
aggregator = FanInAggregator("aggregator")
|
||||
|
||||
# Create workflow with fan-in: executor1 -> [executor2, executor3] -> aggregator
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.set_start_executor(executor1)
|
||||
.add_fan_out_edges(executor1, [executor2, executor3])
|
||||
.add_fan_in_edges([executor2, executor3], aggregator)
|
||||
.build()
|
||||
)
|
||||
|
||||
# Verify build span was created
|
||||
build_spans = [s for s in span_exporter.get_finished_spans() if s.name == "workflow.build"]
|
||||
assert len(build_spans) == 1
|
||||
|
||||
build_span = build_spans[0]
|
||||
assert build_span.attributes is not None
|
||||
assert build_span.attributes.get("workflow.id") == workflow.id
|
||||
assert build_span.attributes.get("workflow.definition") is not None
|
||||
definition = build_span.attributes.get("workflow.definition")
|
||||
assert definition == workflow.model_dump_json()
|
||||
|
||||
# Check build events
|
||||
assert build_span.events is not None
|
||||
build_event_names = [event.name for event in build_span.events]
|
||||
assert "build.started" in build_event_names
|
||||
assert "build.validation_completed" in build_event_names
|
||||
assert "build.completed" in build_event_names
|
||||
|
||||
# Clear spans to separate build from run tracing
|
||||
span_exporter.clear()
|
||||
|
||||
# Run workflow (this should create run spans)
|
||||
events = []
|
||||
async for event in workflow.run_streaming("test input"):
|
||||
events.append(event)
|
||||
|
||||
# Verify workflow executed correctly
|
||||
assert len(executor1.processed_messages) == 1
|
||||
assert executor1.processed_messages[0] == "test input"
|
||||
assert len(executor2.processed_messages) == 1
|
||||
assert executor2.processed_messages[0] == "processed: test input"
|
||||
assert len(executor3.processed_messages) == 1
|
||||
assert executor3.processed_messages[0] == "processed: test input" # executor3 receives from executor1 via fan-out
|
||||
assert len(aggregator.processed_messages) == 1
|
||||
# The aggregator should receive both processed messages from executor2 and executor3
|
||||
aggregated_msg = aggregator.processed_messages[0]
|
||||
assert "second: processed: test input" in aggregated_msg
|
||||
assert "third: processed: test input" in aggregated_msg
|
||||
|
||||
# Check run spans (build spans should not be present after clear)
|
||||
spans = span_exporter.get_finished_spans()
|
||||
|
||||
# Should have workflow span, processing spans, and sending spans
|
||||
workflow_spans = [s for s in spans if s.name == "workflow.run"]
|
||||
processing_spans = [s for s in spans if s.name == "executor.process"]
|
||||
sending_spans = [s for s in spans if s.name == "message.send"]
|
||||
build_spans_after_run = [s for s in spans if s.name == "workflow.build"]
|
||||
|
||||
assert len(workflow_spans) == 1
|
||||
assert len(processing_spans) >= 4 # executor1, executor2, executor3, aggregator
|
||||
assert len(sending_spans) >= 3 # Messages sent between executors
|
||||
assert len(build_spans_after_run) == 0 # No build spans should be present after clear
|
||||
|
||||
# Verify workflow span events
|
||||
workflow_span = workflow_spans[0]
|
||||
assert workflow_span.events is not None
|
||||
event_names = [event.name for event in workflow_span.events]
|
||||
assert "workflow.started" in event_names
|
||||
assert "workflow.completed" in event_names
|
||||
|
||||
# Test fan-in span linking: find the aggregator's processing span
|
||||
aggregator_spans = [s for s in processing_spans if s.attributes and s.attributes.get("executor.id") == "aggregator"]
|
||||
assert len(aggregator_spans) == 1
|
||||
|
||||
aggregator_span = aggregator_spans[0]
|
||||
# The aggregator span should have links to the source spans (from executor2 and executor3)
|
||||
# This tests that FanInEdgeRunner properly handles multiple trace contexts and span IDs
|
||||
assert aggregator_span.links is not None
|
||||
|
||||
# Find the sending spans from executor2 and executor3 by checking parent relationships
|
||||
executor2_processing_spans = [
|
||||
s for s in processing_spans if s.attributes and s.attributes.get("executor.id") == "executor2"
|
||||
]
|
||||
executor3_processing_spans = [
|
||||
s for s in processing_spans if s.attributes and s.attributes.get("executor.id") == "executor3"
|
||||
]
|
||||
|
||||
# Get span IDs from processing spans
|
||||
executor2_processing_span_ids = {format(s.context.span_id, "016x") for s in executor2_processing_spans if s.context}
|
||||
executor3_processing_span_ids = {format(s.context.span_id, "016x") for s in executor3_processing_spans if s.context}
|
||||
|
||||
executor2_sending_spans = [
|
||||
s for s in sending_spans if s.parent and format(s.parent.span_id, "016x") in executor2_processing_span_ids
|
||||
]
|
||||
executor3_sending_spans = [
|
||||
s for s in sending_spans if s.parent and format(s.parent.span_id, "016x") in executor3_processing_span_ids
|
||||
]
|
||||
|
||||
# Verify that we have sending spans from both executors
|
||||
assert len(executor2_sending_spans) >= 1, "Should have at least one sending span from executor2"
|
||||
assert len(executor3_sending_spans) >= 1, "Should have at least one sending span from executor3"
|
||||
|
||||
# Verify that the aggregator span links point to the correct source spans
|
||||
linked_span_ids = {link.context.span_id for link in aggregator_span.links}
|
||||
|
||||
# Should have links from both executor2 and executor3's sending spans
|
||||
executor2_span_ids = {s.context.span_id for s in executor2_sending_spans if s.context}
|
||||
executor3_span_ids = {s.context.span_id for s in executor3_sending_spans if s.context}
|
||||
|
||||
# At least one span from each executor should be linked
|
||||
assert bool(linked_span_ids & executor2_span_ids), "Aggregator should link to executor2's sending span"
|
||||
assert bool(linked_span_ids & executor3_span_ids), "Aggregator should link to executor3's sending span"
|
||||
|
||||
# Should have at least 2 links (one from each source executor)
|
||||
assert len(aggregator_span.links) >= 2, f"Expected at least 2 links, got {len(aggregator_span.links)}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_workflow_error_handling_in_tracing(tracing_enabled: Any, span_exporter: InMemorySpanExporter) -> None:
|
||||
"""Test that workflow errors are properly recorded in traces."""
|
||||
|
||||
class FailingExecutor(Executor):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(id="failing_executor")
|
||||
|
||||
@handler
|
||||
async def handle_message(self, message: str, ctx: WorkflowContext[None]) -> None:
|
||||
raise ValueError("Test error")
|
||||
|
||||
failing_executor = FailingExecutor()
|
||||
workflow = WorkflowBuilder().set_start_executor(failing_executor).build()
|
||||
|
||||
# Run workflow and expect error
|
||||
with pytest.raises(ValueError, match="Test error"):
|
||||
async for _ in workflow.run_streaming("test input"):
|
||||
pass
|
||||
|
||||
spans = span_exporter.get_finished_spans()
|
||||
|
||||
# Find workflow span
|
||||
workflow_spans = [s for s in spans if s.name == "workflow.run"]
|
||||
assert len(workflow_spans) == 1
|
||||
|
||||
workflow_span = workflow_spans[0]
|
||||
|
||||
# Verify error event and status are recorded
|
||||
assert workflow_span.events is not None
|
||||
event_names = [event.name for event in workflow_span.events]
|
||||
assert "workflow.started" in event_names
|
||||
assert "workflow.error" in event_names
|
||||
assert workflow_span.status.status_code.name == "ERROR"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_message_trace_context_serialization() -> None:
|
||||
"""Test that message trace context is properly serialized/deserialized."""
|
||||
ctx = InProcRunnerContext()
|
||||
|
||||
# Create message with trace context
|
||||
message = Message(
|
||||
data="test",
|
||||
source_id="source",
|
||||
target_id="target",
|
||||
trace_contexts=[{"traceparent": "00-trace-span-01"}],
|
||||
source_span_ids=["span123"],
|
||||
)
|
||||
|
||||
await ctx.send_message(message)
|
||||
|
||||
# Get checkpoint state (which serializes messages)
|
||||
state = await ctx.get_checkpoint_state()
|
||||
|
||||
# Check serialized message includes trace context
|
||||
serialized_msg = state["messages"]["source"][0]
|
||||
assert serialized_msg["trace_contexts"] == [{"traceparent": "00-trace-span-01"}]
|
||||
assert serialized_msg["source_span_ids"] == ["span123"]
|
||||
|
||||
# Test deserialization
|
||||
await ctx.set_checkpoint_state(state)
|
||||
restored_messages = await ctx.drain_messages()
|
||||
|
||||
restored_msg = list(restored_messages.values())[0][0]
|
||||
assert restored_msg.trace_context == {"traceparent": "00-trace-span-01"} # Test backward compatibility
|
||||
assert restored_msg.source_span_id == "span123" # Test backward compatibility
|
||||
assert restored_msg.trace_contexts == [{"traceparent": "00-trace-span-01"}] # Test new format
|
||||
assert restored_msg.source_span_ids == ["span123"] # Test new format
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_workflow_build_error_tracing(tracing_enabled: Any, span_exporter: InMemorySpanExporter) -> None:
|
||||
"""Test that build errors are properly recorded in build spans."""
|
||||
|
||||
# Test validation error by not setting start executor
|
||||
builder = WorkflowBuilder()
|
||||
|
||||
with pytest.raises(ValueError, match="Starting executor must be set"):
|
||||
builder.build()
|
||||
|
||||
spans = span_exporter.get_finished_spans()
|
||||
assert len(spans) == 1
|
||||
|
||||
build_span = spans[0]
|
||||
assert build_span.name == "workflow.build"
|
||||
|
||||
# Verify error status and events
|
||||
assert build_span.status.status_code.name == "ERROR"
|
||||
assert build_span.events is not None
|
||||
|
||||
event_names = [event.name for event in build_span.events]
|
||||
assert "build.started" in event_names
|
||||
assert "build.error" in event_names
|
||||
|
||||
# Check error event attributes
|
||||
error_events = [event for event in build_span.events if event.name == "build.error"]
|
||||
assert len(error_events) == 1
|
||||
|
||||
error_event = error_events[0]
|
||||
assert error_event.attributes is not None
|
||||
assert "Starting executor must be set" in str(error_event.attributes.get("build.error.message"))
|
||||
assert error_event.attributes.get("build.error.type") == "ValueError"
|
||||
@@ -288,7 +288,7 @@ async def test_fan_in():
|
||||
def simple_executor() -> Executor:
|
||||
class SimpleExecutor(Executor):
|
||||
@handler
|
||||
async def handle_message(self, message: Message, context: WorkflowContext[None]) -> None:
|
||||
async def handle_message(self, message: str, context: WorkflowContext[None]) -> None:
|
||||
pass
|
||||
|
||||
return SimpleExecutor(id="test_executor")
|
||||
|
||||
Reference in New Issue
Block a user