From 7b94c34c542c702e3897a85f4477bee6e37343cf Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Thu, 6 Aug 2026 23:43:41 -0700 Subject: [PATCH] feat(tracing): per-step obs wrappers inside business Temporal activities (1:1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- src/agentex/lib/core/tracing/trace.py | 70 ++++++++++++--------------- tests/test_temporal_obs_backend.py | 35 +++++++++++++- 2 files changed, 64 insertions(+), 41 deletions(-) diff --git a/src/agentex/lib/core/tracing/trace.py b/src/agentex/lib/core/tracing/trace.py index d3decdb9b..447ed7d9b 100644 --- a/src/agentex/lib/core/tracing/trace.py +++ b/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) diff --git a/tests/test_temporal_obs_backend.py b/tests/test_temporal_obs_backend.py index 34daf1d11..1ca1f5959 100644 --- a/tests/test_temporal_obs_backend.py +++ b/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. + # Option A (tag 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 Option A. 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 (Option A).""" + 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"}