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

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

@@ -28,15 +28,22 @@
 from __future__ import annotations
 
 import os
+import logging
 from typing import Dict, Tuple, Optional
 
-__all__ = ("get_obs_mode", "obs_correlation")
+__all__ = ("get_obs_mode", "obs_correlation", "warn_on_backend_drift")
 
 DD_ONLY = "dd_only"
 LGTM = "lgtm"
 _DEFAULT_MODE = DD_ONLY
 _VALID_MODES = (DD_ONLY, LGTM)
 
+_log = logging.getLogger(__name__)
+# Deduped (expected, actual) drift directions already warned about, so a genuine
+# mismatch logs once instead of once per span. Bounded by construction: at most
+# the 2 direction pairs ("otel"/"ddtrace" either way).
+_WARNED_DRIFT: set[Tuple[str, str]] = set()
+
 
 def get_obs_mode() -> str:
     """Unset/empty/unrecognized -> ``dd_only`` (current behavior)."""
@@ -69,7 +76,7 @@ def _ddtrace_ids() -> Optional[Tuple[str, str]]:
     return None
 
 
-def obs_correlation(prefer_otel: bool = False) -> Dict[str, str]:
+def obs_correlation(expect_otel: bool = False) -> Dict[str, str]:
     """Return ``{"obs_trace_id": ..., "obs_span_id": ...}`` for the active
     observability context, or ``{}`` if none is active.
 
@@ -79,15 +86,15 @@ def obs_correlation(prefer_otel: bool = False) -> Dict[str, str]:
     dotted) keep them addressable via Postgres JSON paths
     (``operation_metadata->>'obs_trace_id'``).
 
-    ``prefer_otel``: on the Temporal path the active span is the temporalio OTel
+    ``expect_otel``: on the Temporal path the active span is the temporalio OTel
     ``TracingInterceptor`` span regardless of ``SGP_OBS_MODE``, so callers there
     read OTel first (falling back to ddtrace) -- otherwise the default ``dd_only``
     mode would read ids for an unrelated ddtrace trace, not the activity span.
 
     Never fabricates ids -- this is a correlation tag, not the span's id.
     """
     try:
-        if prefer_otel:
+        if expect_otel:
             ids = _lgtm_ids() or _ddtrace_ids()
         else:
             ids = _lgtm_ids() if get_obs_mode() == LGTM else _ddtrace_ids()
@@ -97,3 +104,44 @@ def obs_correlation(prefer_otel: bool = False) -> Dict[str, str]:
     if not ids:
         return {}
     return {"obs_trace_id": ids[0], "obs_span_id": ids[1]}
+
+
+def warn_on_backend_drift(expect_otel: bool = False) -> None:
+    """Log once when the EXPECTED obs backend has no active span but the OTHER one
+    does.
+
+    Expected backend = OTel when ``expect_otel`` (the Temporal path, where the
+    interceptor span is OTel regardless of ``SGP_OBS_MODE``), otherwise the backend
+    the mode implies. A mismatch means the mode does not match the tracer actually
+    running at this call site -- e.g. ``dd_only`` configured but the live span is
+    OTel -- which is a real config/instrumentation drift worth surfacing rather
+    than silently correlating against whatever happens to be live.
+
+    Not a hard failure: obs stays fail-open (the caller still reads and falls back,
+    so no correlation is lost). The warning is deduped per direction, so a standing
+    mismatch logs once, not once per span. Probes the expected backend first and
+    returns early when it is live, so the healthy common path never touches the
+    other backend. Never raises."""
+    try:
+        if expect_otel or get_obs_mode() == LGTM:
+            expected, expected_probe, other_probe, actual = "otel", _lgtm_ids, _ddtrace_ids, "ddtrace"
+        else:
+            expected, expected_probe, other_probe, actual = "ddtrace", _ddtrace_ids, _lgtm_ids, "otel"
+        if expected_probe() is not None:
+            return  # expected backend is live -> healthy; skip the other probe
+        if other_probe() is None:
+            return  # nothing live at all -> uninstrumented path, not drift
+        if (expected, actual) not in _WARNED_DRIFT:
+            _WARNED_DRIFT.add((expected, actual))
+            _log.warning(
+                "obs backend drift: expected %s here (SGP_OBS_MODE=%s%s) but the "
+                "active span is %s; correlating against %s. Check SGP_OBS_MODE and "
+                "the running instrumentation.",
+                expected,
+                get_obs_mode(),
+                ", temporal path" if expect_otel else "",
+                actual,
+                actual,
+            )
+    except Exception:  # obs must never fail an app call
+        pass

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

@@ -175,6 +175,7 @@ def open_obs_span(
     name: str,
     business_span_id: Optional[str] = None,
     business_trace_id: Optional[str] = None,
+    expect_otel: bool = False,
 ) -> Optional[ObsSpanHandle]:
     """Open an obs span named ``name`` in the active backend, make it the active
     span, and return a handle carrying its ``{"obs_trace_id","obs_span_id"}``.
@@ -183,6 +184,14 @@ def open_obs_span(
     the reverse tag (``agentex.business_span_id`` / ``agentex.business_trace_id``)
     so you can pivot obs -> business by searching them in Tempo/DD.
 
+    ``expect_otel``: open an OTel wrapper first, regardless of ``SGP_OBS_MODE``.
+    Set on the Temporal path, where the ambient span is the temporalio OTel
+    ``TracingInterceptor`` span regardless of mode -- an OTel wrapper nests under
+    it and yields valid ids, whereas the default ``dd_only`` path would open a
+    ddtrace wrapper, which finds no request context in a worker and returns None
+    (dropping the per-step span and its ids). Falls back to ddtrace if no OTel
+    span materializes.
+
     Returns ``None`` (so the caller falls back to ambient behavior) when the
     backend tracer isn't available or, in ``dd_only``, no request trace is
     active.
@@ -192,8 +201,14 @@ def open_obs_span(
     never fail an app call.
     """
     try:
-        if get_obs_mode() == LGTM:
-            return _open_otel_span(name, business_span_id, business_trace_id)
+        if expect_otel or get_obs_mode() == LGTM:
+            handle = _open_otel_span(name, business_span_id, business_trace_id)
+            if handle is not None or not expect_otel:
+                # In lgtm mode a None handle means "no OTel span -> caller uses the
+                # ambient fallback". Only when expect_otel is set (Temporal path)
+                # do we try ddtrace as a second choice.
+                return handle
+            return _open_ddtrace_span(name, business_span_id, business_trace_id)
         return _open_ddtrace_span(name, business_span_id, business_trace_id)
     except Exception:  # pragma: no cover - backstop; obs must never break a call
         return None
@@ -236,26 +251,27 @@ def _tag_ddtrace_ambient(business_span_id: Optional[str], business_trace_id: Opt
 def tag_ambient_obs_span(
     business_span_id: Optional[str] = None,
     business_trace_id: Optional[str] = None,
-    prefer_otel: bool = False,
+    expect_otel: bool = False,
 ) -> None:
     """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``
-    already made active for this activity and just add
+    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.
 
-    ``prefer_otel``: on the Temporal path the ambient span is the temporalio OTel
+    ``expect_otel``: on the Temporal path the ambient span is the temporalio OTel
     ``TracingInterceptor`` span REGARDLESS of ``SGP_OBS_MODE`` -- so callers there
-    pass ``prefer_otel=True`` to tag OTel first (falling back to ddtrace only if
+    pass ``expect_otel=True`` to tag OTel first (falling back to ddtrace only if
     no valid OTel span is active). Without this, the default ``dd_only`` mode would
     tag an unrelated ddtrace span (or nothing) instead of the real activity span."""
     try:
-        if prefer_otel:
+        if expect_otel:
             if _tag_otel_ambient(business_span_id, business_trace_id):
                 return
             _tag_ddtrace_ambient(business_span_id, business_trace_id)

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

@@ -12,7 +12,7 @@
 from agentex.types.span import Span
 from agentex.lib.utils.logging import make_logger
 from agentex.lib.utils.model_utils import recursive_model_dump
-from agentex.lib.core.tracing.obs_ids import obs_correlation
+from agentex.lib.core.tracing.obs_ids import obs_correlation, warn_on_backend_drift
 from agentex.lib.core.tracing.obs_span import (
     ObsSpanHandle,
     open_obs_span,
@@ -106,35 +106,54 @@ def _run_on_span_end(processor: SyncTracingProcessor, span: Span) -> None:
         )
 
 
+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 by
+    ``TracingActivityName`` (``start-span`` / ``end-span``). Never raises; False
+    when temporalio isn't importable or we're not in an activity."""
+    try:
+        from temporalio import activity
+
+        if not activity.in_activity():
+            return False
+        # Import only AFTER the in_activity() guard: the pure-sync ACP path never
+        # runs this, so it doesn't pull the temporal activities module graph
+        # (activities -> TracingService -> AsyncTracer -> trace, also circular at
+        # import time) into a process that never runs a workflow, and a broken
+        # import can't silently disable the guard on that path. Inside an activity
+        # the graph is fully loaded, so the lazy import is safe -- and it keeps the
+        # discriminator keyed on the enum, not on drifting string literals.
+        # ``activity_type`` round-trips as the enum's str value, which a str-Enum
+        # member compares equal to.
+        from agentex.lib.core.temporal.activities.adk.tracing_activities import (
+            TracingActivityName,
+        )
+
+        return activity.info().activity_type in (
+            TracingActivityName.START_SPAN,
+            TracingActivityName.END_SPAN,
+        )
+    except Exception:
+        return False
+
+
 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.
-    """
+    """True inside ANY Temporal activity. There the ambient span is the temporalio
+    OTel ``TracingInterceptor`` span REGARDLESS of ``SGP_OBS_MODE``, so callers
+    prefer OTel for both the wrapper backend and the correlation read: a plain
+    ``dd_only`` read would target ddtrace, which has no request context in a worker
+    (no inbound HTTP), so ``open_obs_span`` would return None and the fallback ids
+    would be empty -- the business span would persist with no obs_* ids at all.
+    Never raises; False when temporalio isn't importable or we're not in an
+    activity."""
     try:
         from temporalio import activity
 
@@ -148,27 +167,52 @@ 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 ``expect_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).
+
+    Inside ANY activity we also pass ``expect_otel`` to the wrapper and the ambient
+    fallback: the ambient span is the interceptor's OTel span regardless of mode,
+    so a per-step OTel wrapper nests under it and yields valid ids, whereas the
+    default ``dd_only`` path would open a ddtrace wrapper -- which finds no request
+    context in a worker and returns None, leaving the business span with empty
+    obs_* ids.
     """
-    if _in_temporal_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)
-    correlation = handle.correlation if handle is not None else obs_correlation()
+    if _in_tracing_dispatch_activity():
+        warn_on_backend_drift(expect_otel=True)
+        tag_ambient_obs_span(business_span_id=span_id, business_trace_id=trace_id, expect_otel=True)
+        return None, obs_correlation(expect_otel=True)
+    # TODO(obs-followup): two items formerly tracked on the (now-deleted)
+    # _in_temporal_activity docstring, still open after this change:
+    #   (1) TurnTrace RETRY/ASYNC roll-up. A retried business activity now emits a
+    #       full per-step wrapper set PER ATTEMPT, each nested under that attempt's
+    #       RunActivity. Each attempt correlates to the turn on its own, but they
+    #       are not yet rolled up, so a retried turn surfaces as N per-attempt span
+    #       sets rather than one PRIMARY + N RETRY view.
+    #   (2) On a multi-replica worker fleet, assert _OBS_HANDLES stays bounded (no
+    #       leak / OOM) and that obs_trace_id resolves to the turn trace.
+    expect_otel = _in_temporal_activity()
+    warn_on_backend_drift(expect_otel)
+    handle = open_obs_span(
+        name, business_span_id=span_id, business_trace_id=trace_id, expect_otel=expect_otel
+    )
+    correlation = handle.correlation if handle is not None else obs_correlation(expect_otel=expect_otel)
     return handle, correlation
 
 

tests/lib/core/tracing/test_obs_span.py

@@ -351,7 +351,7 @@ def test_lgtm_wrapper_marked_error_when_business_step_raises(self, monkeypatch):
     def test_dd_only_no_ctx_falls_back_to_ambient(self, monkeypatch):
         monkeypatch.setenv("SGP_OBS_MODE", "dd_only")
         _install_fake_ddtrace(monkeypatch, active=False)
-        monkeypatch.setattr("agentex.lib.core.tracing.trace.obs_correlation", lambda: {})
+        monkeypatch.setattr("agentex.lib.core.tracing.trace.obs_correlation", lambda **_k: {})
 
         trace = Trace(processors=[], client=MagicMock(), trace_id="task-run-3")
         span = trace.start_span(name="get_state")