Add BaggageSpanProcessor to populate baggage attributes into OTel spans for Agentic AI cases

Author: mxiamxia

aws-opentelemetry-distro/pyproject.toml

@@ -26,6 +26,7 @@ classifiers = [
 dependencies = [
   "opentelemetry-api == 1.27.0",
   "opentelemetry-sdk == 1.27.0",
+  "opentelemetry-processor-baggage == 0.48b0",
   "opentelemetry-exporter-otlp-proto-grpc == 1.27.0",
   "opentelemetry-exporter-otlp-proto-http == 1.27.0",
   "opentelemetry-propagator-b3 == 1.27.0",

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py

@@ -33,6 +33,7 @@
 from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter as OTLPHttpOTLPMetricExporter
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 from opentelemetry.metrics import set_meter_provider
+from opentelemetry.processor.baggage import ALLOW_ALL_BAGGAGE_KEYS, BaggageSpanProcessor
 from opentelemetry.sdk._configuration import (
     _get_exporter_names,
     _get_id_generator,
@@ -407,6 +408,11 @@ def _customize_span_processors(provider: TracerProvider, resource: Resource) ->
     if _is_lambda_environment():
         provider.add_span_processor(AwsLambdaSpanProcessor())
 
+    # Add baggage attributes to span attributes to support AI Agent use cases
+    # enabling additional data like session ID in spans.
+    if is_agent_observability_enabled():
+        provider.add_span_processor(BaggageSpanProcessor(ALLOW_ALL_BAGGAGE_KEYS))
+
     if not _is_application_signals_enabled():
         return
 

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py

@@ -50,6 +50,7 @@
 from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter
 from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter as OTLPHttpOTLPMetricExporter
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.processor.baggage import BaggageSpanProcessor
 from opentelemetry.sdk.environment_variables import OTEL_TRACES_SAMPLER, OTEL_TRACES_SAMPLER_ARG
 from opentelemetry.sdk.metrics._internal.export import PeriodicExportingMetricReader
 from opentelemetry.sdk.resources import Resource
@@ -340,6 +341,29 @@ def test_customize_span_exporter_with_agent_observability(self):
         os.environ.pop("AGENT_OBSERVABILITY_ENABLED", None)
         os.environ.pop(OTEL_EXPORTER_OTLP_TRACES_ENDPOINT, None)
 
+    def test_customize_span_processors_with_agent_observability(self):
+        mock_tracer_provider: TracerProvider = MagicMock()
+
+        # Test that BaggageSpanProcessor is not added when agent observability is disabled
+        os.environ.pop("AGENT_OBSERVABILITY_ENABLED", None)
+        _customize_span_processors(mock_tracer_provider, Resource.get_empty())
+        self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 0)
+
+        # Reset mock for next test
+        mock_tracer_provider.reset_mock()
+
+        # Test that BaggageSpanProcessor is added when agent observability is enabled
+        os.environ["AGENT_OBSERVABILITY_ENABLED"] = "true"
+        _customize_span_processors(mock_tracer_provider, Resource.get_empty())
+        self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 1)
+
+        # Verify the added processor is BaggageSpanProcessor
+        added_processor = mock_tracer_provider.add_span_processor.call_args_list[0].args[0]
+        self.assertIsInstance(added_processor, BaggageSpanProcessor)
+
+        # Clean up
+        os.environ.pop("AGENT_OBSERVABILITY_ENABLED", None)
+
     def test_customize_span_exporter_sigv4(self):
 
         traces_good_endpoints = [
@@ -574,9 +598,18 @@ def capture_exporter(*args, **kwargs):
 
     def test_customize_span_processors(self):
         mock_tracer_provider: TracerProvider = MagicMock()
+        # Clean up environment to ensure consistent test state
+        os.environ.pop("AGENT_OBSERVABILITY_ENABLED", None)
+        os.environ.pop("OTEL_AWS_APPLICATION_SIGNALS_ENABLED", None)
+        os.environ.pop("OTEL_AWS_APPLICATION_SIGNALS_RUNTIME_ENABLED", None)
+
         _customize_span_processors(mock_tracer_provider, Resource.get_empty())
         self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 0)
 
+        # Reset mock for next test
+        mock_tracer_provider.reset_mock()
+
+        # Test application signals only
         os.environ.setdefault("OTEL_AWS_APPLICATION_SIGNALS_ENABLED", "True")
         os.environ.setdefault("OTEL_AWS_APPLICATION_SIGNALS_RUNTIME_ENABLED", "False")
         _customize_span_processors(mock_tracer_provider, Resource.get_empty())
@@ -586,8 +619,27 @@ def test_customize_span_processors(self):
         second_processor: SpanProcessor = mock_tracer_provider.add_span_processor.call_args_list[1].args[0]
         self.assertIsInstance(second_processor, AwsSpanMetricsProcessor)
 
+        # Reset mock for next test
+        mock_tracer_provider.reset_mock()
+
+        # Test both agent observability and application signals enabled
+        os.environ.setdefault("AGENT_OBSERVABILITY_ENABLED", "true")
+        _customize_span_processors(mock_tracer_provider, Resource.get_empty())
+        self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 3)
+
+        # Verify processors are added in the expected order
+        processors = [call.args[0] for call in mock_tracer_provider.add_span_processor.call_args_list]
+        self.assertIsInstance(processors[0], BaggageSpanProcessor)  # Agent observability processor added first
+        self.assertIsInstance(processors[1], AttributePropagatingSpanProcessor)  # Application signals processors
+        self.assertIsInstance(processors[2], AwsSpanMetricsProcessor)
+
     def test_customize_span_processors_lambda(self):
         mock_tracer_provider: TracerProvider = MagicMock()
+        # Clean up environment to ensure consistent test state
+        os.environ.pop("AGENT_OBSERVABILITY_ENABLED", None)
+        os.environ.pop("OTEL_AWS_APPLICATION_SIGNALS_ENABLED", None)
+        os.environ.pop("AWS_LAMBDA_FUNCTION_NAME", None)
+
         _customize_span_processors(mock_tracer_provider, Resource.get_empty())
         self.assertEqual(mock_tracer_provider.add_span_processor.call_count, 0)