address the PR comment

Author: nikhilc-microsoft

libraries/microsoft-agents-a365-observability-core/microsoft_agents_a365/observability/core/__init__.py

@@ -14,6 +14,7 @@
 from .execution_type import ExecutionType
 from .exporters.enriched_span import EnrichedReadableSpan
 from .exporters.enriching_span_processor import (
+    get_span_enricher,
     register_span_enricher,
     unregister_span_enricher,
 )
@@ -40,6 +41,7 @@
     # Span enrichment
     "register_span_enricher",
     "unregister_span_enricher",
+    "get_span_enricher",
     "EnrichedReadableSpan",
     # Span processor
     "SpanProcessor",

libraries/microsoft-agents-a365-observability-core/microsoft_agents_a365/observability/core/exporters/enriching_span_processor.py

@@ -3,69 +3,84 @@
 
 """Span enrichment support for the Agent365 exporter pipeline."""
 
+import logging
 import threading
 from collections.abc import Callable
 
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 
-# Registry for span enrichers - allows extensions to add attributes to spans before export
-_span_enrichers: list[Callable[[ReadableSpan], ReadableSpan]] = []
-_enrichers_lock = threading.Lock()
+logger = logging.getLogger(__name__)
+
+# Single span enricher - only one platform instrumentor should be active at a time
+_span_enricher: Callable[[ReadableSpan], ReadableSpan] | None = None
+_enricher_lock = threading.Lock()
 
 
 def register_span_enricher(enricher: Callable[[ReadableSpan], ReadableSpan]) -> None:
-    """
-    Register a function that enriches spans before export.
+    """Register the span enricher for the active platform instrumentor.
 
-    Extensions (like Semantic Kernel, LangChain, etc.) can register enrichers
-    that modify spans before they are exported by the BatchSpanProcessor.
+    Only one enricher can be registered at a time since auto-instrumentation
+    is platform-specific (Semantic Kernel, LangChain, or OpenAI Agents).
 
     Args:
-        enricher: A function that takes a ReadableSpan and returns an
-                  enriched ReadableSpan (or the same span if no changes).
+        enricher: Function that takes a ReadableSpan and returns an enriched span.
+
+    Raises:
+        RuntimeError: If an enricher is already registered.
     """
-    with _enrichers_lock:
-        if enricher not in _span_enrichers:
-            _span_enrichers.append(enricher)
+    global _span_enricher
+    with _enricher_lock:
+        if _span_enricher is not None:
+            raise RuntimeError(
+                "A span enricher is already registered. "
+                "Only one platform instrumentor can be active at a time."
+            )
+        _span_enricher = enricher
+        logger.debug("Span enricher registered: %s", enricher.__name__)
 
 
-def unregister_span_enricher(enricher: Callable[[ReadableSpan], ReadableSpan]) -> None:
-    """
-    Remove a previously registered enricher.
+def unregister_span_enricher() -> None:
+    """Unregister the current span enricher.
 
-    Args:
-        enricher: The enricher function to remove.
+    Called during uninstrumentation to clean up.
     """
-    with _enrichers_lock:
-        if enricher in _span_enrichers:
-            _span_enrichers.remove(enricher)
+    global _span_enricher
+    with _enricher_lock:
+        if _span_enricher is not None:
+            logger.debug("Span enricher unregistered: %s", _span_enricher.__name__)
+            _span_enricher = None
 
 
-class _EnrichingBatchSpanProcessor(BatchSpanProcessor):
-    """
-    BatchSpanProcessor that applies registered enrichers before export.
+def get_span_enricher() -> Callable[[ReadableSpan], ReadableSpan] | None:
+    """Get the currently registered span enricher.
 
-    This allows extensions to modify spans after they end but before
-    they are batched and exported.
+    Returns:
+        The registered enricher function, or None if no enricher is registered.
     """
+    with _enricher_lock:
+        return _span_enricher
+
+
+class _EnrichingBatchSpanProcessor(BatchSpanProcessor):
+    """BatchSpanProcessor that applies the registered enricher before batching."""
 
     def on_end(self, span: ReadableSpan) -> None:
-        """
-        Apply all registered enrichers to the span before batching.
+        """Apply the span enricher and pass to parent for batching.
 
         Args:
-            span: The ReadableSpan that has ended.
+            span: The span that has ended.
         """
         enriched_span = span
-        with _enrichers_lock:
-            enrichers = list(_span_enrichers)  # Copy to avoid holding lock during enrichment
 
-        for enricher in enrichers:
+        enricher = get_span_enricher()
+        if enricher is not None:
             try:
-                enriched_span = enricher(enriched_span)
+                enriched_span = enricher(span)
             except Exception:
-                # Don't let enrichment failures break the pipeline
-                pass
+                logger.exception(
+                    "Span enricher %s raised an exception, using original span",
+                    enricher.__name__,
+                )
 
         super().on_end(enriched_span)

libraries/microsoft-agents-a365-observability-extensions-semantickernel/microsoft_agents_a365/observability/extensions/semantickernel/trace_instrumentor.py

@@ -67,7 +67,7 @@ def _uninstrument(self, **kwargs: Any) -> None:
             **kwargs: Optional configuration parameters.
         """
         # Unregister the enricher
-        unregister_span_enricher(enrich_semantic_kernel_span)
+        unregister_span_enricher()
 
         # Shutdown the processor
         if hasattr(self, "_processor"):

tests/observability/core/exporters/test_enriching_span_processor.py

@@ -8,7 +8,7 @@
 
 from microsoft_agents_a365.observability.core.exporters.enriching_span_processor import (
     _EnrichingBatchSpanProcessor,
-    _span_enrichers,
+    get_span_enricher,
     register_span_enricher,
     unregister_span_enricher,
 )
@@ -18,58 +18,66 @@ class TestSpanEnricherRegistry(unittest.TestCase):
     """Test suite for span enricher registration functions."""
 
     def setUp(self):
-        """Clear enrichers before each test."""
-        _span_enrichers.clear()
+        """Ensure clean state before each test."""
+        unregister_span_enricher()
 
     def tearDown(self):
-        """Clear enrichers after each test."""
-        _span_enrichers.clear()
+        """Clean up after each test."""
+        unregister_span_enricher()
 
     def test_register_and_unregister_enricher(self):
-        """Test that enrichers can be registered and unregistered."""
+        """Test that enricher can be registered and unregistered."""
 
-        # Define a simple enricher
         def my_enricher(span):
             return span
 
-        # Register
-        register_span_enricher(my_enricher)
-        self.assertIn(my_enricher, _span_enrichers)
-        self.assertEqual(len(_span_enrichers), 1)
+        # Initially no enricher
+        self.assertIsNone(get_span_enricher())
 
-        # Duplicate registration should not add again
+        # Register
         register_span_enricher(my_enricher)
-        self.assertEqual(len(_span_enrichers), 1)
+        self.assertEqual(get_span_enricher(), my_enricher)
 
         # Unregister
-        unregister_span_enricher(my_enricher)
-        self.assertNotIn(my_enricher, _span_enrichers)
-        self.assertEqual(len(_span_enrichers), 0)
+        unregister_span_enricher()
+        self.assertIsNone(get_span_enricher())
 
-    def test_unregister_nonexistent_enricher_does_not_raise(self):
-        """Test that unregistering a non-existent enricher doesn't raise an error."""
+    def test_register_second_enricher_raises_error(self):
+        """Test that registering a second enricher raises RuntimeError."""
 
-        def my_enricher(span):
+        def enricher_one(span):
             return span
 
+        def enricher_two(span):
+            return span
+
+        register_span_enricher(enricher_one)
+
+        with self.assertRaises(RuntimeError) as context:
+            register_span_enricher(enricher_two)
+
+        self.assertIn("already registered", str(context.exception))
+
+    def test_unregister_when_none_registered_is_safe(self):
+        """Test that unregistering when no enricher is registered doesn't raise."""
         # Should not raise
-        unregister_span_enricher(my_enricher)
-        self.assertEqual(len(_span_enrichers), 0)
+        unregister_span_enricher()
+        self.assertIsNone(get_span_enricher())
 
 
 class TestEnrichingBatchSpanProcessor(unittest.TestCase):
     """Test suite for _EnrichingBatchSpanProcessor."""
 
     def setUp(self):
-        """Clear enrichers before each test."""
-        _span_enrichers.clear()
+        """Ensure clean state before each test."""
+        unregister_span_enricher()
 
     def tearDown(self):
-        """Clear enrichers after each test."""
-        _span_enrichers.clear()
+        """Clean up after each test."""
+        unregister_span_enricher()
 
-    def test_on_end_applies_enrichers_to_span(self):
-        """Test that on_end applies all registered enrichers to the span."""
+    def test_on_end_applies_enricher_to_span(self):
+        """Test that on_end applies the registered enricher to the span."""
         # Create processor with a mock exporter
         mock_exporter = Mock()
         processor = _EnrichingBatchSpanProcessor(mock_exporter)
@@ -103,23 +111,14 @@ def enricher(span):
         processor.shutdown()
 
     def test_on_end_continues_if_enricher_raises_exception(self):
-        """Test that on_end continues processing even if an enricher raises an exception."""
+        """Test that on_end continues processing even if enricher raises an exception."""
         mock_exporter = Mock()
         processor = _EnrichingBatchSpanProcessor(mock_exporter)
 
-        # Track which enrichers were called
-        called_enrichers = []
-
         def failing_enricher(span):
-            called_enrichers.append("failing")
             raise ValueError("Enricher failed!")
 
-        def succeeding_enricher(span):
-            called_enrichers.append("succeeding")
-            return span
-
         register_span_enricher(failing_enricher)
-        register_span_enricher(succeeding_enricher)
 
         # Create a mock span
         original_span = Mock(name="original_span")
@@ -130,9 +129,22 @@ def succeeding_enricher(span):
         # Should not raise despite failing enricher
         processor.on_end(original_span)
 
-        # Verify both enrichers were called (failing one didn't stop the chain)
-        self.assertIn("failing", called_enrichers)
-        self.assertIn("succeeding", called_enrichers)
+        # Cleanup
+        processor.shutdown()
+
+    def test_on_end_works_without_enricher(self):
+        """Test that on_end works when no enricher is registered."""
+        mock_exporter = Mock()
+        processor = _EnrichingBatchSpanProcessor(mock_exporter)
+
+        # Create a mock span (no enricher registered)
+        original_span = Mock(name="original_span")
+        original_span.context = Mock()
+        original_span.context.trace_id = 123
+        original_span.context.span_id = 456
+
+        # Should not raise
+        processor.on_end(original_span)
 
         # Cleanup
         processor.shutdown()