Add suppress_invoke_agent_input option to OpenAI trace processor

Copilot · sergioescalera · Copilot · commit ddb5dc6ffa22 · 2026-01-09T18:12:45.000Z
Co-authored-by: sergioescalera &lt;8428450+sergioescalera@users.noreply.github.com&gt;
diff --git a/libraries/microsoft-agents-a365-observability-extensions-openai/microsoft_agents_a365/observability/extensions/openai/trace_instrumentor.py b/libraries/microsoft-agents-a365-observability-extensions-openai/microsoft_agents_a365/observability/extensions/openai/trace_instrumentor.py
@@ -47,6 +47,7 @@ def _instrument(self, **kwargs: Any) -> None:
         """Instruments the OpenAI Agents SDK with Microsoft Agent 365 tracing."""
         tracer_name = kwargs["tracer_name"] if kwargs.get("tracer_name") else None
         tracer_version = kwargs["tracer_version"] if kwargs.get("tracer_version") else None
+        suppress_invoke_agent_input = kwargs.get("suppress_invoke_agent_input", False)
 
         # Get the configured Microsoft Agent 365 Tracer
         try:
@@ -64,7 +65,9 @@ def _instrument(self, **kwargs: Any) -> None:
 
         agent365_tracer = cast(Tracer, tracer)
 
-        set_trace_processors([OpenAIAgentsTraceProcessor(agent365_tracer)])
+        set_trace_processors(
+            [OpenAIAgentsTraceProcessor(agent365_tracer, suppress_invoke_agent_input)]
+        )
 
     def _uninstrument(self, **kwargs: Any) -> None:
         pass
diff --git a/libraries/microsoft-agents-a365-observability-extensions-openai/microsoft_agents_a365/observability/extensions/openai/trace_processor.py b/libraries/microsoft-agents-a365-observability-extensions-openai/microsoft_agents_a365/observability/extensions/openai/trace_processor.py
@@ -69,8 +69,9 @@
 class OpenAIAgentsTraceProcessor(TracingProcessor):
     _MAX_HANDOFFS_IN_FLIGHT = 1000
 
-    def __init__(self, tracer: Tracer) -> None:
+    def __init__(self, tracer: Tracer, suppress_invoke_agent_input: bool = False) -> None:
         self._tracer = tracer
+        self._suppress_invoke_agent_input = suppress_invoke_agent_input
         self._root_spans: dict[str, OtelSpan] = {}
         self._otel_spans: dict[str, OtelSpan] = {}
         self._tokens: dict[str, object] = {}
@@ -79,6 +80,8 @@ def __init__(self, tracer: Tracer) -> None:
         # 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:
@@ -89,6 +92,11 @@ 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 _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 on_trace_start(self, trace: Trace) -> None:
         """Called when a trace is started.
 
@@ -134,6 +142,12 @@ 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.
 
@@ -154,7 +168,11 @@ def on_span_end(self, span: Span[Any]) -> None:
                 otel_span.set_attribute(GEN_AI_OUTPUT_MESSAGES_KEY, response.model_dump_json())
                 for k, v in get_attributes_from_response(response):
                     otel_span.set_attribute(k, v)
-            if hasattr(data, "input") and (input := data.input):
+            # Only record input messages if not suppressing or not in InvokeAgent scope
+            should_suppress = self._suppress_invoke_agent_input and self._is_in_invoke_agent_scope(
+                span.trace_id
+            )
+            if not should_suppress 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):
@@ -164,8 +182,18 @@ def on_span_end(self, span: Span[Any]) -> None:
                 elif TYPE_CHECKING:
                     assert_never(input)
         elif isinstance(data, GenerationSpanData):
-            for k, v in get_attributes_from_generation_span_data(data):
-                otel_span.set_attribute(k, v)
+            # Only record input messages if not suppressing or not in InvokeAgent scope
+            should_suppress = self._suppress_invoke_agent_input and self._is_in_invoke_agent_scope(
+                span.trace_id
+            )
+            if not should_suppress:
+                for k, v in get_attributes_from_generation_span_data(data):
+                    otel_span.set_attribute(k, v)
+            else:
+                # Still set attributes other than input messages
+                for k, v in get_attributes_from_generation_span_data(data):
+                    if k != GEN_AI_INPUT_MESSAGES_KEY:
+                        otel_span.set_attribute(k, v)
             self._stamp_custom_parent(otel_span, span.trace_id)
             otel_span.update_name(
                 f"{otel_span.attributes[GEN_AI_OPERATION_NAME_KEY]} {otel_span.attributes[GEN_AI_REQUEST_MODEL_KEY]}"
@@ -194,6 +222,12 @@ 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:
diff --git a/tests/observability/extensions/openai/test_prompt_suppression.py b/tests/observability/extensions/openai/test_prompt_suppression.py
@@ -0,0 +1,283 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import unittest
+from datetime import datetime
+from unittest.mock import Mock
+
+from agents.tracing import Span
+from agents.tracing.span_data import AgentSpanData, GenerationSpanData, ResponseSpanData
+from microsoft_agents_a365.observability.core import configure, get_tracer
+from microsoft_agents_a365.observability.core.constants import GEN_AI_INPUT_MESSAGES_KEY
+from microsoft_agents_a365.observability.extensions.openai.trace_processor import (
+    OpenAIAgentsTraceProcessor,
+)
+from openai.types.responses import Response
+
+
+class TestPromptSuppression(unittest.TestCase):
+    """Unit tests for prompt suppression functionality in OpenAIAgentsTraceProcessor."""
+
+    @classmethod
+    def setUpClass(cls):
+        """Set up test environment once for all tests."""
+        configure(
+            service_name="test-service-prompt-suppression",
+            service_namespace="test-namespace-prompt-suppression",
+        )
+
+    def setUp(self):
+        """Set up each test with a fresh processor and mock tracer."""
+        self.tracer = get_tracer()
+        self.mock_otel_span = Mock()
+        self.mock_otel_span.attributes = {}
+        self.mock_otel_span.get_span_context.return_value = Mock(
+            trace_id="test-trace-id", span_id="test-span-id"
+        )
+        
+        # Track attributes set on the span
+        def set_attribute_side_effect(key, value):
+            self.mock_otel_span.attributes[key] = value
+        
+        self.mock_otel_span.set_attribute = Mock(side_effect=set_attribute_side_effect)
+        self.mock_otel_span.update_name = Mock()
+        self.mock_otel_span.set_status = Mock()
+        self.mock_otel_span.end = Mock()
+
+        # Mock the tracer's start_span method
+        self.original_start_span = self.tracer.start_span
+        self.tracer.start_span = Mock(return_value=self.mock_otel_span)
+
+    def tearDown(self):
+        """Clean up after each test."""
+        self.tracer.start_span = self.original_start_span
+
+    def test_does_not_record_input_messages_when_suppression_enabled_in_agent_scope(self):
+        """Test that input messages are not recorded when suppression is enabled and in agent scope."""
+        processor = OpenAIAgentsTraceProcessor(self.tracer, suppress_invoke_agent_input=True)
+
+        trace_id = "trace-suppress"
+        now = datetime.now().isoformat()
+
+        # Start an agent span to create InvokeAgent scope
+        agent_span = Mock(spec=Span)
+        agent_span.span_id = "agent-span"
+        agent_span.trace_id = trace_id
+        agent_span.parent_id = None
+        agent_span.started_at = now
+        agent_span.ended_at = None
+        agent_span.span_data = AgentSpanData(name="TestAgent")
+
+        processor.on_span_start(agent_span)
+
+        # Now create a generation span with input (proper format - list of message dicts)
+        gen_span = Mock(spec=Span)
+        gen_span.span_id = "gen-span"
+        gen_span.trace_id = trace_id
+        gen_span.parent_id = "agent-span"
+        gen_span.started_at = now
+        gen_span.ended_at = now
+        gen_span.span_data = GenerationSpanData(
+            model="gpt-4",
+            input=[{"role": "user", "content": "Hello prompt"}]
+        )
+
+        processor.on_span_start(gen_span)
+        processor.on_span_end(gen_span)
+
+        # Verify that set_attribute was called but NOT with GEN_AI_INPUT_MESSAGES_KEY
+        attribute_keys = [call[0][0] for call in self.mock_otel_span.set_attribute.call_args_list]
+        self.assertNotIn(
+            GEN_AI_INPUT_MESSAGES_KEY,
+            attribute_keys,
+            "GEN_AI_INPUT_MESSAGES_KEY should not be set when suppression is enabled",
+        )
+
+    def test_records_input_messages_when_suppression_disabled(self):
+        """Test that input messages are recorded when suppression is disabled (default)."""
+        processor = OpenAIAgentsTraceProcessor(self.tracer, suppress_invoke_agent_input=False)
+
+        trace_id = "trace-allow"
+        now = datetime.now().isoformat()
+
+        # Start an agent span
+        agent_span = Mock(spec=Span)
+        agent_span.span_id = "agent-span-2"
+        agent_span.trace_id = trace_id
+        agent_span.parent_id = None
+        agent_span.started_at = now
+        agent_span.ended_at = None
+        agent_span.span_data = AgentSpanData(name="TestAgent")
+
+        processor.on_span_start(agent_span)
+
+        # Create a generation span with input (proper format - list of message dicts)
+        gen_span = Mock(spec=Span)
+        gen_span.span_id = "gen-span-2"
+        gen_span.trace_id = trace_id
+        gen_span.parent_id = "agent-span-2"
+        gen_span.started_at = now
+        gen_span.ended_at = now
+        gen_span.span_data = GenerationSpanData(
+            model="gpt-4",
+            input=[{"role": "user", "content": "Hello prompt"}]
+        )
+
+        processor.on_span_start(gen_span)
+        processor.on_span_end(gen_span)
+
+        # Verify that set_attribute was called with GEN_AI_INPUT_MESSAGES_KEY
+        attribute_keys = [call[0][0] for call in self.mock_otel_span.set_attribute.call_args_list]
+        self.assertIn(
+            GEN_AI_INPUT_MESSAGES_KEY,
+            attribute_keys,
+            "GEN_AI_INPUT_MESSAGES_KEY should be set when suppression is disabled",
+        )
+
+    def test_suppresses_input_on_response_spans_when_enabled(self):
+        """Test that input is suppressed on response spans when suppression is enabled."""
+        processor = OpenAIAgentsTraceProcessor(self.tracer, suppress_invoke_agent_input=True)
+
+        trace_id = "trace-resp"
+        now = datetime.now().isoformat()
+
+        # Start an agent span
+        agent_span = Mock(spec=Span)
+        agent_span.span_id = "agent-span-3"
+        agent_span.trace_id = trace_id
+        agent_span.parent_id = None
+        agent_span.started_at = now
+        agent_span.ended_at = None
+        agent_span.span_data = AgentSpanData(name="TestAgent")
+
+        processor.on_span_start(agent_span)
+
+        # Create a response span with input
+        resp_span = Mock(spec=Span)
+        resp_span.span_id = "resp-span"
+        resp_span.trace_id = trace_id
+        resp_span.parent_id = "agent-span-3"
+        resp_span.started_at = now
+        resp_span.ended_at = now
+
+        # Create mock response data with all required attributes
+        mock_response = Mock(spec=Response)
+        mock_response.model_dump_json.return_value = '{"output": "test"}'
+        mock_response.tools = None
+        mock_response.usage = None
+        mock_response.output = None
+        mock_response.instructions = None
+        mock_response.model = "gpt-4"
+        mock_response.model_dump.return_value = {}
+
+        resp_span.span_data = Mock(spec=ResponseSpanData)
+        resp_span.span_data.response = mock_response
+        resp_span.span_data.input = "Prompt text"
+
+        processor.on_span_start(resp_span)
+        processor.on_span_end(resp_span)
+
+        # Verify that set_attribute was called but NOT with GEN_AI_INPUT_MESSAGES_KEY for input
+        attribute_keys = [call[0][0] for call in self.mock_otel_span.set_attribute.call_args_list]
+        self.assertNotIn(
+            GEN_AI_INPUT_MESSAGES_KEY,
+            attribute_keys,
+            "GEN_AI_INPUT_MESSAGES_KEY should not be set for response span when suppression is enabled",
+        )
+
+    def test_records_input_outside_agent_scope_even_when_suppression_enabled(self):
+        """Test that input messages are recorded outside agent scope even when suppression is enabled."""
+        processor = OpenAIAgentsTraceProcessor(self.tracer, suppress_invoke_agent_input=True)
+
+        trace_id = "trace-outside"
+        now = datetime.now().isoformat()
+
+        # Create a generation span WITHOUT an agent span (outside InvokeAgent scope)
+        gen_span = Mock(spec=Span)
+        gen_span.span_id = "gen-span-outside"
+        gen_span.trace_id = trace_id
+        gen_span.parent_id = None
+        gen_span.started_at = now
+        gen_span.ended_at = now
+        gen_span.span_data = GenerationSpanData(
+            model="gpt-4",
+            input=[{"role": "user", "content": "Hello prompt"}]
+        )
+
+        processor.on_span_start(gen_span)
+        processor.on_span_end(gen_span)
+
+        # Verify that set_attribute WAS called with GEN_AI_INPUT_MESSAGES_KEY
+        # because we're not in an InvokeAgent scope
+        attribute_keys = [call[0][0] for call in self.mock_otel_span.set_attribute.call_args_list]
+        self.assertIn(
+            GEN_AI_INPUT_MESSAGES_KEY,
+            attribute_keys,
+            "GEN_AI_INPUT_MESSAGES_KEY should be set when outside InvokeAgent scope",
+        )
+
+    def test_default_suppression_is_false(self):
+        """Test that the default value for suppress_invoke_agent_input is False."""
+        processor = OpenAIAgentsTraceProcessor(self.tracer)
+
+        self.assertFalse(
+            processor._suppress_invoke_agent_input,
+            "Default value for suppress_invoke_agent_input should be False",
+        )
+
+    def test_agent_span_tracking_cleanup(self):
+        """Test that agent span tracking is properly cleaned up when spans end."""
+        processor = OpenAIAgentsTraceProcessor(self.tracer, suppress_invoke_agent_input=True)
+
+        trace_id = "trace-cleanup"
+        now = datetime.now().isoformat()
+
+        # Start an agent span
+        agent_span = Mock(spec=Span)
+        agent_span.span_id = "agent-span-cleanup"
+        agent_span.trace_id = trace_id
+        agent_span.parent_id = None
+        agent_span.started_at = now
+        agent_span.ended_at = now
+        agent_span.span_data = AgentSpanData(name="TestAgent")
+
+        processor.on_span_start(agent_span)
+
+        # Verify the span is tracked
+        self.assertIn(trace_id, processor._active_agent_spans)
+        self.assertIn(agent_span.span_id, processor._active_agent_spans[trace_id])
+
+        # End the agent span
+        processor.on_span_end(agent_span)
+
+        # Verify the tracking is cleaned up
+        self.assertNotIn(trace_id, processor._active_agent_spans)
+
+
+def run_tests():
+    """Run all prompt suppression tests."""
+    print("🧪 Running prompt suppression tests...")
+    print("=" * 80)
+
+    loader = unittest.TestLoader()
+    suite = loader.loadTestsFromTestCase(TestPromptSuppression)
+
+    runner = unittest.TextTestRunner(verbosity=2)
+    result = runner.run(suite)
+
+    print("\n" + "=" * 80)
+    print("🏁 Test Summary:")
+    print(f"Tests run: {result.testsRun}")
+    print(f"Failures: {len(result.failures)}")
+    print(f"Errors: {len(result.errors)}")
+
+    if result.wasSuccessful():
+        print("🎉 All tests passed!")
+        return True
+    else:
+        print("🔧 Some tests failed. Check output above.")
+        return False
+
+
+if __name__ == "__main__":
+    success = run_tests()
+    exit(0 if success else 1)
