feat(tracing): per-step obs wrappers inside business Temporal activities (1:1)

Previously _begin_obs skipped the obs wrapper for ANY Temporal activity (Option
A) and only stamped the ambient RunActivity span, so all business spans in a turn
collapsed onto ONE obs span (52:1). But inside a *business* activity, start_span
and end_span run in the SAME process, so a wrapper is safe there. Option A is
only required for the SDK's own dispatched START_SPAN/END_SPAN activities (the
in_temporal_workflow path), where start and end are separate activities on
possibly different workers.

Discriminate on activity type: _in_tracing_dispatch_activity() is true only for
the "start-span"/"end-span" activities. For everything else (sync, or a business
activity) open a real per-step wrapper — it nests under the interceptor's ambient
RunActivity span and closes in-process, giving each business span its own obs
span (1:1), matching the sync path. The bounded _OBS_HANDLES registry backstops
any mis-discrimination.

Author: NiteshDhanpal

src/agentex/lib/core/tracing/obs_span.py

@@ -241,10 +241,10 @@ def tag_ambient_obs_span(
     """Stamp the reverse tag onto the CURRENTLY ACTIVE obs span -- without opening
     a new one.
 
-    Used on the Temporal path (see ``trace._in_temporal_activity``): there we must
-    NOT open our own wrapper span, because start_span/end_span run as separate
-    activities on possibly different workers and the wrapper could never be
-    closed. Instead we lean on the span the Temporal OTel ``TracingInterceptor``
+    Used inside the SDK's dispatched start-span/end-span activities (see
+    ``trace._in_tracing_dispatch_activity``): there we must NOT open our own
+    wrapper span, because start_span/end_span run as separate activities on
+    possibly different workers and the wrapper could never be closed. Instead we lean on the span the Temporal OTel ``TracingInterceptor``
     already made active for this activity and just add
     ``agentex.business_span_id`` / ``agentex.business_trace_id`` so the obs -> business
     pivot still works. Best-effort; never raises.

src/agentex/lib/core/tracing/trace.py

@@ -106,39 +106,25 @@ def _run_on_span_end(processor: SyncTracingProcessor, span: Span) -> None:
         )
 
 
-def _in_temporal_activity() -> bool:
-    """True when executing inside a Temporal activity.
-
-    On the Temporal path ``start_span`` and ``end_span`` run as SEPARATE
-    activities (START_SPAN / END_SPAN) that Temporal can route to DIFFERENT
-    worker processes. A wrapper obs span opened in the START_SPAN activity could
-    therefore never be closed by END_SPAN -- its handle lives in another
-    process's ``_OBS_HANDLES`` -- so it would leak (unbounded, OOM risk) and its
-    persisted ``obs_span_id`` would dangle (the span is never .end()ed, so never
-    exported to Tempo).
-
-    So inside an activity we do NOT open our own wrapper. We lean on the span the
-    Temporal OTel ``TracingInterceptor`` (see ``core/tracing/temporal.py`` +
-    scale-agentex-python#485) already made active for this activity -- which is
-    rooted under the turn's propagated trace -- and merely stamp the reverse tag
-    onto it (``tag_ambient_obs_span``). That keeps trace-level correlation with
-    no cross-process handle to leak.
-
-    Never raises; returns False when temporalio isn't importable.
-
-    TODO(obs-followup): this intentionally drops the *named per-step* wrapper on
-    the Temporal path (obs_span_id becomes the ambient activity span, not a
-    step-named span) and does NOT add TurnTrace RETRY/ASYNC roll-up -- retried
-    turns still surface as N unlinked spans. Follow-up diff should (a) optionally
-    materialize a self-contained named wrapper inside a single activity using the
-    span's own start/end timestamps, and (b) build the TurnTrace roll-up.
-    Test-later: on a multi-replica worker fleet, assert _OBS_HANDLES stays
-    bounded (no leak / OOM) and that obs_trace_id resolves to the turn trace.
-    """
+def _in_tracing_dispatch_activity() -> bool:
+    """True only when running inside the SDK's OWN dispatched START_SPAN / END_SPAN
+    activity (the ``in_temporal_workflow()`` path, where a workflow runs span start
+    and end as SEPARATE activities that Temporal can route to different workers).
+
+    That is the one case a per-step obs wrapper can't work: the wrapper opened in
+    the START_SPAN activity could never be closed by the END_SPAN activity. A span
+    created directly inside a *business* activity (an agent turn's own
+    ``adk.tracing.span``) runs start AND end in the same activity process, so a
+    wrapper there is safe -- it nests under the interceptor's ambient RunActivity
+    span and closes in-process. The tracing dispatch activities are named
+    ``start-span`` / ``end-span`` (``TracingActivityName``). Never raises; False
+    when temporalio isn't importable or we're not in an activity."""
     try:
         from temporalio import activity
 
-        return activity.in_activity()
+        if not activity.in_activity():
+            return False
+        return activity.info().activity_type in ("start-span", "end-span")
     except Exception:
         return False
 
@@ -148,23 +134,27 @@ def _begin_obs(
     span_id: str,
     trace_id: str | None,
 ) -> tuple[ObsSpanHandle | None, dict[str, str]]:
-    """Open the obs wrapper for a business span (or, inside a Temporal activity,
-    tag the ambient interceptor span) and return ``(handle, correlation)``.
+    """Open the obs wrapper for a business span and return ``(handle, correlation)``.
 
     Shared by ``Trace.start_span`` and ``AsyncTrace.start_span`` so the two paths
     can't drift. The wrapper is named for the step so ``obs_span_id`` is
     stable/meaningful (not an arbitrary innermost httpx span), and it carries the
     reverse tag (business span/trace id) for the obs -> business pivot.
 
-    Temporal path: we do NOT open our own wrapper -- start_span / end_span run as
-    separate activities on possibly different workers, so the handle could never
-    be closed. Instead we tag the span the temporalio OTel ``TracingInterceptor``
-    already made active. That span is OTel REGARDLESS of ``SGP_OBS_MODE``, so we
-    pass ``prefer_otel=True`` to both the tag and the correlation read -- otherwise
-    the default ``dd_only`` mode would tag/read an unrelated ddtrace span and the
-    ids would point at the wrong trace. See ``_in_temporal_activity``.
+    We open a real per-step wrapper on the sync path AND inside a *business*
+    Temporal activity -- there the wrapper nests under the interceptor's ambient
+    RunActivity span and start/end run in-process, so it closes cleanly and each
+    business step gets its own obs span (1:1), just like sync.
+
+    The ONE exception is the SDK's own dispatched START_SPAN / END_SPAN activity
+    (a workflow calling ``adk.tracing`` -- see ``_in_tracing_dispatch_activity``):
+    there start and end are separate activities on possibly different workers, so
+    a wrapper could never be closed. We fall back to tagging the ambient
+    interceptor span instead, with ``prefer_otel=True`` (the interceptor span is
+    OTel regardless of ``SGP_OBS_MODE``, so a plain ``dd_only`` read would
+    otherwise point at an unrelated ddtrace span).
     """
-    if _in_temporal_activity():
+    if _in_tracing_dispatch_activity():
         tag_ambient_obs_span(business_span_id=span_id, business_trace_id=trace_id, prefer_otel=True)
         return None, obs_correlation(prefer_otel=True)
     handle = open_obs_span(name, business_span_id=span_id, business_trace_id=trace_id)

tests/test_temporal_obs_backend.py

@@ -67,8 +67,10 @@ def _activate_otel_span(monkeypatch: pytest.MonkeyPatch) -> _RecordingOtelSpan:
 
 def test_temporal_path_tags_and_reads_otel_in_dd_only(monkeypatch: pytest.MonkeyPatch) -> None:
     # Default/dd_only mode is exactly where the old code went to ddtrace.
+    # Tagging the ambient interceptor span (no wrapper) now applies only inside
+    # the SDK's dispatched START_SPAN/END_SPAN activity, not any activity.
     monkeypatch.setenv("SGP_OBS_MODE", "dd_only")
-    monkeypatch.setattr(trace_mod, "_in_temporal_activity", lambda: True)
+    monkeypatch.setattr(trace_mod, "_in_tracing_dispatch_activity", lambda: True)
     activity_span = _activate_otel_span(monkeypatch)
 
     trace_obj = Trace(processors=[], client=cast(Any, object()), trace_id="trace-1")
@@ -132,3 +134,34 @@ def current_span(self) -> _FakeDDSpan:
     assert tagged["agentex.business_trace_id"] == "bt"
     # The invalid OTel span was NOT tagged.
     assert invalid.attributes == {}
+
+
+class _FakeHandle:
+    def __init__(self, corr):
+        self.correlation = corr
+
+
+def test_begin_obs_opens_wrapper_outside_dispatch_activity(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Sync path or inside a business Temporal activity: open a per-step wrapper
+    (1:1), not the ambient-span tag. Each business span gets its own obs span."""
+    monkeypatch.setattr(trace_mod, "_in_tracing_dispatch_activity", lambda: False)
+    monkeypatch.setattr(
+        trace_mod, "open_obs_span",
+        lambda *a, **k: _FakeHandle({"obs_trace_id": "t1", "obs_span_id": "s1"}),
+    )
+    handle, corr = trace_mod._begin_obs("mortgage.classify_intent", "bs", "bt")
+    assert handle is not None
+    assert corr == {"obs_trace_id": "t1", "obs_span_id": "s1"}
+
+
+def test_begin_obs_tags_ambient_inside_dispatch_activity(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Inside the dispatched START_SPAN/END_SPAN activity: no wrapper (would leak
+    across activities); tag the ambient interceptor span instead."""
+    monkeypatch.setattr(trace_mod, "_in_tracing_dispatch_activity", lambda: True)
+    tagged: dict = {}
+    monkeypatch.setattr(trace_mod, "tag_ambient_obs_span", lambda **k: tagged.update(k))
+    monkeypatch.setattr(trace_mod, "obs_correlation", lambda **k: {"obs_trace_id": "amb", "obs_span_id": "amb"})
+    handle, corr = trace_mod._begin_obs("mortgage.advisor.turn", "bs", "bt")
+    assert handle is None
+    assert tagged.get("business_span_id") == "bs" and tagged.get("prefer_otel") is True
+    assert corr == {"obs_trace_id": "amb", "obs_span_id": "amb"}