From f105e9007f2a1f251b95680fe58276a11680de50 Mon Sep 17 00:00:00 2001 From: Tim Conley Date: Wed, 5 Aug 2026 14:14:27 -0700 Subject: [PATCH] Fix OpenInference trace patch lifecycle --- .../openai_agents/_temporal_openai_agents.py | 71 ++++++++++++++----- .../openai_agents/test_openai_tracing.py | 25 +++++++ 2 files changed, 80 insertions(+), 16 deletions(-) diff --git a/temporalio/contrib/openai_agents/_temporal_openai_agents.py b/temporalio/contrib/openai_agents/_temporal_openai_agents.py index 43594657f..ffe851a78 100644 --- a/temporalio/contrib/openai_agents/_temporal_openai_agents.py +++ b/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() diff --git a/tests/contrib/openai_agents/test_openai_tracing.py b/tests/contrib/openai_agents/test_openai_tracing.py index facc3212b..8fa195bb1 100644 --- a/tests/contrib/openai_agents/test_openai_tracing.py +++ b/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)