Fix OpenInference trace patch lifecycle

temporalio/contrib/openai_agents/_temporal_openai_agents.py

@@ -2,6 +2,7 @@
 
 import dataclasses
 import json
+import threading
 import typing
 from collections.abc import AsyncIterator, Callable, Iterator, Sequence
 from contextlib import asynccontextmanager, contextmanager
@@ -54,6 +55,55 @@
     )
 
 
+_otel_trace_start_patch_lock = threading.RLock()
+_otel_trace_start_patch_ref_count = 0
+_otel_trace_start_original: Callable[..., typing.Any] | None = None
+
+
+def _install_otel_trace_start_patch() -> None:
+    """Make an OpenInference root span current while tracing is configured."""
+    global _otel_trace_start_original
+    global _otel_trace_start_patch_ref_count
+
+    from openinference.instrumentation.openai_agents._processor import (
+        OpenInferenceTracingProcessor,
+    )
+    from opentelemetry.context import attach
+    from opentelemetry.trace import set_span_in_context
+
+    with _otel_trace_start_patch_lock:
+        if _otel_trace_start_patch_ref_count == 0:
+            _otel_trace_start_original = OpenInferenceTracingProcessor.on_trace_start
+
+            def on_trace_start(self: typing.Any, trace: Trace) -> None:  # type: ignore[reportUnusedFunction]
+                _otel_trace_start_original(self, trace)  # type: ignore[operator]
+                attach(set_span_in_context(self._root_spans[trace.trace_id]))
+
+            setattr(OpenInferenceTracingProcessor, "on_trace_start", on_trace_start)
+        _otel_trace_start_patch_ref_count += 1
+
+
+def _uninstall_otel_trace_start_patch() -> None:
+    """Restore OpenInference after the final tracing context exits."""
+    global _otel_trace_start_original
+    global _otel_trace_start_patch_ref_count
+
+    from openinference.instrumentation.openai_agents._processor import (
+        OpenInferenceTracingProcessor,
+    )
+
+    with _otel_trace_start_patch_lock:
+        _otel_trace_start_patch_ref_count -= 1
+        if _otel_trace_start_patch_ref_count == 0:
+            if _otel_trace_start_original is not None:
+                setattr(
+                    OpenInferenceTracingProcessor,
+                    "on_trace_start",
+                    _otel_trace_start_original,
+                )
+            _otel_trace_start_original = None
+
+
 @contextmanager
 def _set_open_ai_agent_temporal_overrides(
     model_params: ModelActivityParameters,
@@ -377,28 +427,15 @@ def tracing_context(self) -> Iterator[None]:
         """
         # Set up OTEL instrumentation if exporters are provided
         otel_instrumentor = None
+        otel_trace_start_patch_installed = False
         if self._use_otel_instrumentation and not self._instrumented:
             from openinference.instrumentation.openai_agents import (
                 OpenAIAgentsInstrumentor,
             )
-            from openinference.instrumentation.openai_agents._processor import (
-                OpenInferenceTracingProcessor,
-            )
             from opentelemetry import trace
-            from opentelemetry.context import attach
-            from opentelemetry.trace import set_span_in_context
-
-            # Unfortunate monkey patching is needed to ensure the trace is set in context so we can propagate it.
-            original_on_trace_start = OpenInferenceTracingProcessor.on_trace_start
-
-            def on_trace_start(self, trace: Trace) -> None:  # type: ignore[reportMissingParameterType]
-                original_on_trace_start(self, trace)
-                otel_span = self._root_spans[trace.trace_id]
-                attach(set_span_in_context(otel_span))
-
-            OpenInferenceTracingProcessor.on_trace_start = on_trace_start  # type:ignore[method-assign]
 
-            # Set up instrumentor
+            _install_otel_trace_start_patch()
+            otel_trace_start_patch_installed = True
             otel_instrumentor = OpenAIAgentsInstrumentor()
             otel_instrumentor.instrument(tracer_provider=trace.get_tracer_provider())
             self._instrumented = True
@@ -408,3 +445,5 @@ def on_trace_start(self, trace: Trace) -> None:  # type: ignore[reportMissingPar
             # Clean up OTEL instrumentation
             if otel_instrumentor is not None:
                 otel_instrumentor.uninstrument()
+            if otel_trace_start_patch_installed:
+                _uninstall_otel_trace_start_patch()

tests/contrib/openai_agents/test_openai_tracing.py

@@ -11,6 +11,7 @@
 
 from temporalio import activity, workflow
 from temporalio.client import Client
+from temporalio.contrib.openai_agents import _temporal_openai_agents
 from temporalio.contrib.openai_agents.testing import (
     AgentEnvironment,
 )
@@ -50,6 +51,30 @@ def force_flush(self) -> None:
         pass
 
 
+def test_otel_trace_start_patch_does_not_nest() -> None:
+    from openinference.instrumentation.openai_agents._processor import (
+        OpenInferenceTracingProcessor,
+    )
+
+    original = OpenInferenceTracingProcessor.on_trace_start
+    _temporal_openai_agents._install_otel_trace_start_patch()
+    try:
+        installed_patch = OpenInferenceTracingProcessor.on_trace_start
+        assert installed_patch is not original
+
+        _temporal_openai_agents._install_otel_trace_start_patch()
+        try:
+            assert OpenInferenceTracingProcessor.on_trace_start is installed_patch
+        finally:
+            _temporal_openai_agents._uninstall_otel_trace_start_patch()
+
+        assert OpenInferenceTracingProcessor.on_trace_start is installed_patch
+    finally:
+        _temporal_openai_agents._uninstall_otel_trace_start_patch()
+
+    assert OpenInferenceTracingProcessor.on_trace_start is original
+
+
 async def test_tracing(client: Client):
     async with AgentEnvironment(model=research_mock_model()) as env:
         client = env.applied_on_client(client)