1212from agentex .types .span import Span
1313from agentex .lib .utils .logging import make_logger
1414from agentex .lib .utils .model_utils import recursive_model_dump
15- from agentex .lib .core .tracing .obs_ids import obs_correlation
15+ from agentex .lib .core .tracing .obs_ids import obs_correlation , warn_on_backend_drift
1616from agentex .lib .core .tracing .obs_span import (
1717 ObsSpanHandle ,
1818 open_obs_span ,
@@ -106,35 +106,54 @@ def _run_on_span_end(processor: SyncTracingProcessor, span: Span) -> None:
106106 )
107107
108108
109+ def _in_tracing_dispatch_activity () -> bool :
110+ """True only when running inside the SDK's OWN dispatched START_SPAN / END_SPAN
111+ activity (the ``in_temporal_workflow()`` path, where a workflow runs span start
112+ and end as SEPARATE activities that Temporal can route to different workers).
113+
114+ That is the one case a per-step obs wrapper can't work: the wrapper opened in
115+ the START_SPAN activity could never be closed by the END_SPAN activity. A span
116+ created directly inside a *business* activity (an agent turn's own
117+ ``adk.tracing.span``) runs start AND end in the same activity process, so a
118+ wrapper there is safe -- it nests under the interceptor's ambient RunActivity
119+ span and closes in-process. The tracing dispatch activities are named by
120+ ``TracingActivityName`` (``start-span`` / ``end-span``). Never raises; False
121+ when temporalio isn't importable or we're not in an activity."""
122+ try :
123+ from temporalio import activity
124+
125+ if not activity .in_activity ():
126+ return False
127+ # Import only AFTER the in_activity() guard: the pure-sync ACP path never
128+ # runs this, so it doesn't pull the temporal activities module graph
129+ # (activities -> TracingService -> AsyncTracer -> trace, also circular at
130+ # import time) into a process that never runs a workflow, and a broken
131+ # import can't silently disable the guard on that path. Inside an activity
132+ # the graph is fully loaded, so the lazy import is safe -- and it keeps the
133+ # discriminator keyed on the enum, not on drifting string literals.
134+ # ``activity_type`` round-trips as the enum's str value, which a str-Enum
135+ # member compares equal to.
136+ from agentex .lib .core .temporal .activities .adk .tracing_activities import (
137+ TracingActivityName ,
138+ )
139+
140+ return activity .info ().activity_type in (
141+ TracingActivityName .START_SPAN ,
142+ TracingActivityName .END_SPAN ,
143+ )
144+ except Exception :
145+ return False
146+
147+
109148def _in_temporal_activity () -> bool :
110- """True when executing inside a Temporal activity.
111-
112- On the Temporal path ``start_span`` and ``end_span`` run as SEPARATE
113- activities (START_SPAN / END_SPAN) that Temporal can route to DIFFERENT
114- worker processes. A wrapper obs span opened in the START_SPAN activity could
115- therefore never be closed by END_SPAN -- its handle lives in another
116- process's ``_OBS_HANDLES`` -- so it would leak (unbounded, OOM risk) and its
117- persisted ``obs_span_id`` would dangle (the span is never .end()ed, so never
118- exported to Tempo).
119-
120- So inside an activity we do NOT open our own wrapper. We lean on the span the
121- Temporal OTel ``TracingInterceptor`` (see ``core/tracing/temporal.py`` +
122- scale-agentex-python#485) already made active for this activity -- which is
123- rooted under the turn's propagated trace -- and merely stamp the reverse tag
124- onto it (``tag_ambient_obs_span``). That keeps trace-level correlation with
125- no cross-process handle to leak.
126-
127- Never raises; returns False when temporalio isn't importable.
128-
129- TODO(obs-followup): this intentionally drops the *named per-step* wrapper on
130- the Temporal path (obs_span_id becomes the ambient activity span, not a
131- step-named span) and does NOT add TurnTrace RETRY/ASYNC roll-up -- retried
132- turns still surface as N unlinked spans. Follow-up diff should (a) optionally
133- materialize a self-contained named wrapper inside a single activity using the
134- span's own start/end timestamps, and (b) build the TurnTrace roll-up.
135- Test-later: on a multi-replica worker fleet, assert _OBS_HANDLES stays
136- bounded (no leak / OOM) and that obs_trace_id resolves to the turn trace.
137- """
149+ """True inside ANY Temporal activity. There the ambient span is the temporalio
150+ OTel ``TracingInterceptor`` span REGARDLESS of ``SGP_OBS_MODE``, so callers
151+ prefer OTel for both the wrapper backend and the correlation read: a plain
152+ ``dd_only`` read would target ddtrace, which has no request context in a worker
153+ (no inbound HTTP), so ``open_obs_span`` would return None and the fallback ids
154+ would be empty -- the business span would persist with no obs_* ids at all.
155+ Never raises; False when temporalio isn't importable or we're not in an
156+ activity."""
138157 try :
139158 from temporalio import activity
140159
@@ -148,27 +167,52 @@ def _begin_obs(
148167 span_id : str ,
149168 trace_id : str | None ,
150169) -> tuple [ObsSpanHandle | None , dict [str , str ]]:
151- """Open the obs wrapper for a business span (or, inside a Temporal activity,
152- tag the ambient interceptor span) and return ``(handle, correlation)``.
170+ """Open the obs wrapper for a business span and return ``(handle, correlation)``.
153171
154172 Shared by ``Trace.start_span`` and ``AsyncTrace.start_span`` so the two paths
155173 can't drift. The wrapper is named for the step so ``obs_span_id`` is
156174 stable/meaningful (not an arbitrary innermost httpx span), and it carries the
157175 reverse tag (business span/trace id) for the obs -> business pivot.
158176
159- Temporal path: we do NOT open our own wrapper -- start_span / end_span run as
160- separate activities on possibly different workers, so the handle could never
161- be closed. Instead we tag the span the temporalio OTel ``TracingInterceptor``
162- already made active. That span is OTel REGARDLESS of ``SGP_OBS_MODE``, so we
163- pass ``prefer_otel=True`` to both the tag and the correlation read -- otherwise
164- the default ``dd_only`` mode would tag/read an unrelated ddtrace span and the
165- ids would point at the wrong trace. See ``_in_temporal_activity``.
177+ We open a real per-step wrapper on the sync path AND inside a *business*
178+ Temporal activity -- there the wrapper nests under the interceptor's ambient
179+ RunActivity span and start/end run in-process, so it closes cleanly and each
180+ business step gets its own obs span (1:1), just like sync.
181+
182+ The ONE exception is the SDK's own dispatched START_SPAN / END_SPAN activity
183+ (a workflow calling ``adk.tracing`` -- see ``_in_tracing_dispatch_activity``):
184+ there start and end are separate activities on possibly different workers, so
185+ a wrapper could never be closed. We fall back to tagging the ambient
186+ interceptor span instead, with ``expect_otel=True`` (the interceptor span is
187+ OTel regardless of ``SGP_OBS_MODE``, so a plain ``dd_only`` read would
188+ otherwise point at an unrelated ddtrace span).
189+
190+ Inside ANY activity we also pass ``expect_otel`` to the wrapper and the ambient
191+ fallback: the ambient span is the interceptor's OTel span regardless of mode,
192+ so a per-step OTel wrapper nests under it and yields valid ids, whereas the
193+ default ``dd_only`` path would open a ddtrace wrapper -- which finds no request
194+ context in a worker and returns None, leaving the business span with empty
195+ obs_* ids.
166196 """
167- if _in_temporal_activity ():
168- tag_ambient_obs_span (business_span_id = span_id , business_trace_id = trace_id , prefer_otel = True )
169- return None , obs_correlation (prefer_otel = True )
170- handle = open_obs_span (name , business_span_id = span_id , business_trace_id = trace_id )
171- correlation = handle .correlation if handle is not None else obs_correlation ()
197+ if _in_tracing_dispatch_activity ():
198+ warn_on_backend_drift (expect_otel = True )
199+ tag_ambient_obs_span (business_span_id = span_id , business_trace_id = trace_id , expect_otel = True )
200+ return None , obs_correlation (expect_otel = True )
201+ # TODO(obs-followup): two items formerly tracked on the (now-deleted)
202+ # _in_temporal_activity docstring, still open after this change:
203+ # (1) TurnTrace RETRY/ASYNC roll-up. A retried business activity now emits a
204+ # full per-step wrapper set PER ATTEMPT, each nested under that attempt's
205+ # RunActivity. Each attempt correlates to the turn on its own, but they
206+ # are not yet rolled up, so a retried turn surfaces as N per-attempt span
207+ # sets rather than one PRIMARY + N RETRY view.
208+ # (2) On a multi-replica worker fleet, assert _OBS_HANDLES stays bounded (no
209+ # leak / OOM) and that obs_trace_id resolves to the turn trace.
210+ expect_otel = _in_temporal_activity ()
211+ warn_on_backend_drift (expect_otel )
212+ handle = open_obs_span (
213+ name , business_span_id = span_id , business_trace_id = trace_id , expect_otel = expect_otel
214+ )
215+ correlation = handle .correlation if handle is not None else obs_correlation (expect_otel = expect_otel )
172216 return handle , correlation
173217
174218
0 commit comments