Refactor: Use parent span operation name check instead of tracking dictionary

Co-authored-by: sergioescalera <8428450+sergioescalera@users.noreply.github.com>

Author: Copilot

libraries/microsoft-agents-a365-observability-extensions-openai/microsoft_agents_a365/observability/extensions/openai/trace_processor.py

@@ -80,8 +80,6 @@ def __init__(self, tracer: Tracer, suppress_invoke_agent_input: bool = False) ->
         # Use an OrderedDict and _MAX_HANDOFFS_IN_FLIGHT to cap the size of the dict
         # in case there are large numbers of orphaned handoffs
         self._reverse_handoffs_dict: OrderedDict[str, str] = OrderedDict()
-        # Track active agent spans per trace to determine if we're in an InvokeAgent scope
-        self._active_agent_spans: dict[str, set[str]] = {}
 
     # helper
     def _stamp_custom_parent(self, otel_span: OtelSpan, trace_id: str) -> None:
@@ -92,14 +90,27 @@ def _stamp_custom_parent(self, otel_span: OtelSpan, trace_id: str) -> None:
         pid_hex = "0x" + ot_trace.format_span_id(sc.span_id)
         otel_span.set_attribute(CUSTOM_PARENT_SPAN_ID_KEY, pid_hex)
 
-    def _should_suppress_input(self, trace_id: str) -> bool:
-        """Check if input messages should be suppressed for the given trace."""
-        return self._suppress_invoke_agent_input and self._is_in_invoke_agent_scope(trace_id)
-
-    def _is_in_invoke_agent_scope(self, trace_id: str) -> bool:
-        """Check if we're currently inside an InvokeAgent scope for the given trace."""
-        agent_spans = self._active_agent_spans.get(trace_id, set())
-        return len(agent_spans) > 0
+    def _should_suppress_input(self, span: Span[Any]) -> bool:
+        """Check if input messages should be suppressed for the given span.
+        
+        Args:
+            span: The span to check.
+            
+        Returns:
+            True if suppression is enabled and the span has an InvokeAgent parent.
+        """
+        if not self._suppress_invoke_agent_input:
+            return False
+        
+        # Check if the parent span is an InvokeAgent span by looking at its operation name
+        if span.parent_id:
+            parent_otel_span = self._otel_spans.get(span.parent_id)
+            if parent_otel_span and hasattr(parent_otel_span, 'attributes'):
+                operation_name = parent_otel_span.attributes.get(GEN_AI_OPERATION_NAME_KEY)
+                if operation_name == INVOKE_AGENT_OPERATION_NAME:
+                    return True
+        
+        return False
 
     def on_trace_start(self, trace: Trace) -> None:
         """Called when a trace is started.
@@ -146,12 +157,6 @@ def on_span_start(self, span: Span[Any]) -> None:
         self._otel_spans[span.span_id] = otel_span
         self._tokens[span.span_id] = attach(set_span_in_context(otel_span))
 
-        # Track agent spans for InvokeAgent scope detection
-        if isinstance(span.span_data, AgentSpanData):
-            if span.trace_id not in self._active_agent_spans:
-                self._active_agent_spans[span.trace_id] = set()
-            self._active_agent_spans[span.trace_id].add(span.span_id)
-
     def on_span_end(self, span: Span[Any]) -> None:
         """Called when a span is finished. Should not block or raise exceptions.
 
@@ -173,7 +178,7 @@ def on_span_end(self, span: Span[Any]) -> None:
                 for k, v in get_attributes_from_response(response):
                     otel_span.set_attribute(k, v)
             # Only record input messages if not suppressing or not in InvokeAgent scope
-            if not self._should_suppress_input(span.trace_id) and hasattr(data, "input") and (input := data.input):
+            if not self._should_suppress_input(span) and hasattr(data, "input") and (input := data.input):
                 if isinstance(input, str):
                     otel_span.set_attribute(GEN_AI_INPUT_MESSAGES_KEY, input)
                 elif isinstance(input, list):
@@ -184,7 +189,7 @@ def on_span_end(self, span: Span[Any]) -> None:
                     assert_never(input)
         elif isinstance(data, GenerationSpanData):
             # Collect all attributes once and filter if suppression is enabled
-            should_suppress = self._should_suppress_input(span.trace_id)
+            should_suppress = self._should_suppress_input(span)
             for k, v in get_attributes_from_generation_span_data(data):
                 # Skip input messages if suppression is enabled and in InvokeAgent scope
                 if should_suppress and k == GEN_AI_INPUT_MESSAGES_KEY:
@@ -218,12 +223,6 @@ def on_span_end(self, span: Span[Any]) -> None:
                 otel_span.set_attribute(GEN_AI_GRAPH_NODE_PARENT_ID, parent_node)
             otel_span.update_name(f"{INVOKE_AGENT_OPERATION_NAME} {get_span_name(span)}")
 
-            # Clean up agent span tracking
-            if span.trace_id in self._active_agent_spans:
-                self._active_agent_spans[span.trace_id].discard(span.span_id)
-                if not self._active_agent_spans[span.trace_id]:
-                    del self._active_agent_spans[span.trace_id]
-
         end_time: int | None = None
         if span.ended_at:
             try:

tests/observability/extensions/openai/test_prompt_suppression.py

@@ -49,45 +49,18 @@ def test_can_disable_suppression(self):
             "suppress_invoke_agent_input should be False when explicitly set",
         )
 
-    def test_has_active_agent_spans_tracking(self):
-        """Test that the processor has the required tracking data structure."""
+    def test_has_should_suppress_input_method(self):
+        """Test that the processor has the helper method for suppression logic."""
         tracer = get_tracer()
         processor = OpenAIAgentsTraceProcessor(tracer, suppress_invoke_agent_input=True)
 
         self.assertTrue(
-            hasattr(processor, "_active_agent_spans"),
-            "Processor should have _active_agent_spans attribute",
-        )
-        self.assertIsInstance(
-            processor._active_agent_spans,
-            dict,
-            "_active_agent_spans should be a dictionary",
-        )
-
-    def test_has_is_in_invoke_agent_scope_method(self):
-        """Test that the processor has the helper method for scope detection."""
-        tracer = get_tracer()
-        processor = OpenAIAgentsTraceProcessor(tracer, suppress_invoke_agent_input=True)
-
-        self.assertTrue(
-            hasattr(processor, "_is_in_invoke_agent_scope"),
-            "Processor should have _is_in_invoke_agent_scope method",
+            hasattr(processor, "_should_suppress_input"),
+            "Processor should have _should_suppress_input method",
         )
         self.assertTrue(
-            callable(processor._is_in_invoke_agent_scope),
-            "_is_in_invoke_agent_scope should be callable",
-        )
-
-    def test_is_in_invoke_agent_scope_returns_false_for_empty_trace(self):
-        """Test that _is_in_invoke_agent_scope returns False for unknown trace."""
-        tracer = get_tracer()
-        processor = OpenAIAgentsTraceProcessor(tracer, suppress_invoke_agent_input=True)
-
-        result = processor._is_in_invoke_agent_scope("unknown-trace-id")
-
-        self.assertFalse(
-            result,
-            "_is_in_invoke_agent_scope should return False for traces with no active agent spans",
+            callable(processor._should_suppress_input),
+            "_should_suppress_input should be callable",
         )