From 99f8d25c1eca53c3c87b32ab9d33fe4badcbe73e Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Sat, 1 Aug 2026 10:58:30 -0700 Subject: [PATCH 01/12] feat(tracing): correlate business spans with obs via dedicated wrapper span Model the business<->observability correlation edge on the emit side, with no schema migration (rides the existing operation_metadata JSONB). - obs_ids: standardize the correlation keys to obs_trace_id/obs_span_id (underscored, JSON-path friendly); remove the non-working `dual` mode (it required an in-process ddtrace<->OTel bridge that can't exist -- you can't run ddtrace-run and the OTel operator together, and DD_TRACE_OTEL_ENABLED is a single tracer). `dual` now safely degrades to dd_only. Harden obs_correlation to never raise. - obs_span (new): when the SDK creates a business span it opens a dedicated obs span named for that step and makes it active, so obs_span_id is stable and meaningful (a named span with its httpx call nested underneath) instead of an arbitrary innermost instrumentation span. Backends: OTel in lgtm; ddtrace in dd_only but only when a request trace is already active (avoids orphan root traces in un-instrumented agents). Reverse tag: stamps agentex.business_span_id / agentex.business_trace_id onto the obs span so the pivot is bidirectional. - trace: wire the wrapper into start_span/end_span (sync + async). Observability can never fail an app call -- every path is guarded and is a no-op when the tracer isn't configured. - tests: obs_ids (mode degrade + keys), obs_span (both backends, non-interference, never-fails, reverse tag), and the 3-turn mortgage Turn-2 example pinned as an executable contract. Co-Authored-By: Claude Opus 4.8 --- src/agentex/lib/core/tracing/obs_ids.py | 33 +- src/agentex/lib/core/tracing/obs_span.py | 153 +++++++++ src/agentex/lib/core/tracing/trace.py | 51 ++- tests/lib/core/tracing/test_obs_ids.py | 127 ++++++++ tests/lib/core/tracing/test_obs_span.py | 399 +++++++++++++++++++++++ 5 files changed, 740 insertions(+), 23 deletions(-) create mode 100644 src/agentex/lib/core/tracing/obs_span.py create mode 100644 tests/lib/core/tracing/test_obs_ids.py create mode 100644 tests/lib/core/tracing/test_obs_span.py diff --git a/src/agentex/lib/core/tracing/obs_ids.py b/src/agentex/lib/core/tracing/obs_ids.py index 99c6b2555..bac500b20 100644 --- a/src/agentex/lib/core/tracing/obs_ids.py +++ b/src/agentex/lib/core/tracing/obs_ids.py @@ -11,11 +11,16 @@ persisted business span to the Tempo/Datadog trace for the turn that produced it, while the business trace still groups the entire run by task id. -Source selection follows SGP_OBS_MODE, matching egp-api-backend: +Source selection follows SGP_OBS_MODE: - unset / "dd_only": ddtrace context (current stack) - - "dual": OTel/LGTM preferred, ddtrace fallback - "lgtm": OTel/LGTM only +("dual" was removed: co-resident ddtrace+OTel can't be bridged in-process -- +you can't run ddtrace-run and the OTel operator's auto-instrumentation in the +same process, and DD_TRACE_OTEL_ENABLED yields a single tracer with nothing to +bridge. Two-backend export is a collector fan-out under "lgtm", not a mode here. +An unrecognized SGP_OBS_MODE -- including a stale "dual" -- degrades to dd_only.) + This never fabricates ids -- if no observability context is active, it returns an empty dict and the span is simply not tagged. """ @@ -27,10 +32,9 @@ __all__ = ("get_obs_mode", "obs_correlation") DD_ONLY = "dd_only" -DUAL = "dual" LGTM = "lgtm" _DEFAULT_MODE = DD_ONLY -_VALID_MODES = (DD_ONLY, DUAL, LGTM) +_VALID_MODES = (DD_ONLY, LGTM) def get_obs_mode() -> str: @@ -65,19 +69,22 @@ def _ddtrace_ids() -> Optional[Tuple[str, str]]: def obs_correlation() -> Dict[str, str]: - """Return ``{"obs.trace_id": ..., "obs.span_id": ...}`` for the active + """Return ``{"obs_trace_id": ..., "obs_span_id": ...}`` for the active observability context, or ``{}`` if none is active. + These land in the business span's ``data`` -> egp ``operation_metadata`` + (an existing JSONB column, GIN-indexed) -> ClickHouse ``metadata_raw``, so + the correlation edge needs no schema migration. Underscored keys (not + dotted) keep them addressable via Postgres JSON paths + (``operation_metadata->>'obs_trace_id'``). + Never fabricates ids -- this is a correlation tag, not the span's id. """ - mode = get_obs_mode() - if mode == LGTM: - ids = _lgtm_ids() - elif mode == DUAL: - ids = _lgtm_ids() or _ddtrace_ids() - else: # dd_only - ids = _ddtrace_ids() + try: + ids = _lgtm_ids() if get_obs_mode() == LGTM else _ddtrace_ids() + except Exception: # obs must never fail an app call + return {} if not ids: return {} - return {"obs.trace_id": ids[0], "obs.span_id": ids[1]} + return {"obs_trace_id": ids[0], "obs_span_id": ids[1]} diff --git a/src/agentex/lib/core/tracing/obs_span.py b/src/agentex/lib/core/tracing/obs_span.py new file mode 100644 index 000000000..77186d2ad --- /dev/null +++ b/src/agentex/lib/core/tracing/obs_span.py @@ -0,0 +1,153 @@ +"""Dedicated per-business-span observability wrapper span. + +Capturing obs ids from "whatever instrumentation span happens to be innermost +at emit time" is coarse -- it could be an arbitrary httpx-client span, and every +business span in a request would collapse onto the same request/activity span. + +Instead, when the SDK creates a business span we open a **real obs span named +for that step and make it active**. Then: + - ``obs_span_id`` is stable and meaningful (a span named for the business + step, not an arbitrary leaf), and + - any nested instrumentation (httpx, db, ...) parents under it. + +The wrapper's own trace_id/span_id are read directly from its span context, so +the correlation tag is deterministic regardless of what else is on the stack. + +Backend follows ``SGP_OBS_MODE``: + - ``lgtm`` -> an OpenTelemetry span (the convergence target). + - ``dd_only`` -> a ddtrace span, but ONLY when a ddtrace trace is already + active for the request. Opening one unconditionally would emit orphan root + traces in un-instrumented (bare-uvicorn, no ddtrace-run) agents, so when + nothing is active we return ``None`` and the caller keeps its ambient + behavior. + +No-op when the relevant tracer isn't importable. Never raises -- observability +must never break a business span. +""" +from __future__ import annotations + +from typing import Callable, Dict, Optional + +from agentex.lib.core.tracing.obs_ids import LGTM, get_obs_mode + +__all__ = ("ObsSpanHandle", "open_obs_span", "close_obs_span") + +# Instrumentation scope name so these wrapper spans are identifiable in Tempo/DD. +_TRACER_NAME = "agentex.business" + +# Reverse-tag attribute keys: the business span/trace ids stamped onto the obs +# span so you can pivot obs -> business (search these in Tempo/DD). +_ATTR_BUSINESS_SPAN_ID = "agentex.business_span_id" +_ATTR_BUSINESS_TRACE_ID = "agentex.business_trace_id" + + +class ObsSpanHandle: + """Live handle for an open wrapper span: the correlation tag read from it + plus a backend-specific closer (detach/end or finish).""" + + __slots__ = ("correlation", "_close") + + def __init__(self, correlation: Dict[str, str], close: Callable[[], None]): + self.correlation = correlation + self._close = close + + +def _hex_ids(trace_id: int, span_id: int) -> Dict[str, str]: + """W3C-hex form: 32-hex trace, 16-hex span.""" + return { + "obs_trace_id": format(trace_id, "032x"), + "obs_span_id": format(span_id, "016x"), + } + + +def _open_otel_span( + name: str, + business_span_id: Optional[str], + business_trace_id: Optional[str], +) -> Optional[ObsSpanHandle]: + try: + from opentelemetry import context, trace + except ImportError: + return None + try: + span = trace.get_tracer(_TRACER_NAME).start_span(name) + if business_span_id: + span.set_attribute(_ATTR_BUSINESS_SPAN_ID, business_span_id) + if business_trace_id: + span.set_attribute(_ATTR_BUSINESS_TRACE_ID, business_trace_id) + token = context.attach(trace.set_span_in_context(span)) + sc = span.get_span_context() + correlation = _hex_ids(sc.trace_id, sc.span_id) if (sc and sc.is_valid) else {} + + def _close() -> None: + try: + context.detach(token) + finally: + span.end() + + return ObsSpanHandle(correlation, _close) + except Exception: # pragma: no cover - best-effort; never break the business span + return None + + +def _open_ddtrace_span( + name: str, + business_span_id: Optional[str], + business_trace_id: Optional[str], +) -> Optional[ObsSpanHandle]: + try: + from ddtrace.trace import tracer + except ImportError: + return None + try: + # Only wrap when ddtrace is actually tracing the request; otherwise a + # wrapper would be an orphan root trace in an un-instrumented process. + if tracer.current_trace_context() is None: + return None + span = tracer.start_span(name, activate=True) + if business_span_id: + span.set_tag(_ATTR_BUSINESS_SPAN_ID, business_span_id) + if business_trace_id: + span.set_tag(_ATTR_BUSINESS_TRACE_ID, business_trace_id) + correlation = _hex_ids(span.trace_id, span.span_id) if span.trace_id else {} + return ObsSpanHandle(correlation, span.finish) + except Exception: # pragma: no cover - best-effort + return None + + +def open_obs_span( + name: str, + business_span_id: Optional[str] = None, + business_trace_id: Optional[str] = None, +) -> 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"}``. + + ``business_span_id`` / ``business_trace_id`` are stamped onto the obs span as + the reverse tag (``agentex.business_span_id`` / ``agentex.business_trace_id``) + so you can pivot obs -> business by searching them in Tempo/DD. + + 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. + + Never raises: a top-level guard backstops anything the backend helpers + don't (e.g. a broken tracer install raising on import) so observability can + never fail an app call. + """ + try: + if get_obs_mode() == LGTM: + return _open_otel_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 + + +def close_obs_span(handle: Optional[ObsSpanHandle]) -> None: + """Close the wrapper span (detach + end, or finish). Safe on ``None``.""" + if handle is None: + return + try: + handle._close() + except Exception: # pragma: no cover - best-effort + pass diff --git a/src/agentex/lib/core/tracing/trace.py b/src/agentex/lib/core/tracing/trace.py index c3ec91bc3..031e69e9c 100644 --- a/src/agentex/lib/core/tracing/trace.py +++ b/src/agentex/lib/core/tracing/trace.py @@ -12,6 +12,11 @@ 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_span import ( + ObsSpanHandle, + close_obs_span, + open_obs_span, +) from agentex.lib.core.tracing.span_error import set_span_error from agentex.lib.core.tracing.span_queue import ( SpanEventType, @@ -49,6 +54,8 @@ def __init__( self.processors = processors self.client = client self.trace_id = trace_id + # Live per-business-span obs wrapper spans, keyed by business span id. + self._obs_handles: dict[str, ObsSpanHandle] = {} def start_span( self, @@ -80,13 +87,19 @@ def start_span( serialized_input = recursive_model_dump(input) if input else None serialized_data = recursive_model_dump(data) if data else None - # Tag the business span with the active observability trace_id/span_id - # (OTel/ddtrace) so it can be correlated to the per-turn obs trace. The - # business trace_id stays the run-level task id -- see obs_ids.py. - obs = obs_correlation() + # Open a dedicated obs wrapper span named for this step and make it + # active, so obs_span_id is stable/meaningful (not an arbitrary innermost + # httpx span). It also carries the reverse tag (business span/trace id) + # so you can pivot obs -> business in Tempo/DD. Falls back to the ambient + # obs context (ddtrace) when not in lgtm mode. Business trace_id stays the + # run-level task id. + id = str(uuid.uuid4()) + obs_handle = open_obs_span( + name, business_span_id=id, business_trace_id=self.trace_id + ) + obs = obs_handle.correlation if obs_handle is not None else obs_correlation() if obs: serialized_data = {**(serialized_data or {}), **obs} - id = str(uuid.uuid4()) span = Span( id=id, @@ -98,6 +111,8 @@ def start_span( data=serialized_data, task_id=task_id, ) + if obs_handle is not None: + self._obs_handles[span.id] = obs_handle for processor in self.processors: processor.on_span_start(span) @@ -120,6 +135,9 @@ def end_span( if span.end_time is None: span.end_time = datetime.now(UTC) + # Close the dedicated obs wrapper span (detach context + end it). + close_obs_span(self._obs_handles.pop(span.id, None)) + span.input = recursive_model_dump(span.input) if span.input else None span.output = recursive_model_dump(span.output) if span.output else None span.data = recursive_model_dump(span.data) if span.data else None @@ -206,6 +224,8 @@ def __init__( self.client = client self.trace_id = trace_id self._span_queue = span_queue or get_default_span_queue() + # Live per-business-span obs wrapper spans, keyed by business span id. + self._obs_handles: dict[str, ObsSpanHandle] = {} async def start_span( self, @@ -236,13 +256,19 @@ async def start_span( serialized_input = recursive_model_dump(input) if input else None serialized_data = recursive_model_dump(data) if data else None - # Tag the business span with the active observability trace_id/span_id - # (OTel/ddtrace) so it can be correlated to the per-turn obs trace. The - # business trace_id stays the run-level task id -- see obs_ids.py. - obs = obs_correlation() + # Open a dedicated obs wrapper span named for this step and make it + # active, so obs_span_id is stable/meaningful (not an arbitrary innermost + # httpx span). It also carries the reverse tag (business span/trace id) + # so you can pivot obs -> business in Tempo/DD. Falls back to the ambient + # obs context (ddtrace) when not in lgtm mode. Business trace_id stays the + # run-level task id. + id = str(uuid.uuid4()) + obs_handle = open_obs_span( + name, business_span_id=id, business_trace_id=self.trace_id + ) + obs = obs_handle.correlation if obs_handle is not None else obs_correlation() if obs: serialized_data = {**(serialized_data or {}), **obs} - id = str(uuid.uuid4()) span = Span( id=id, @@ -254,6 +280,8 @@ async def start_span( data=serialized_data, task_id=task_id, ) + if obs_handle is not None: + self._obs_handles[span.id] = obs_handle if self.processors: self._span_queue.enqueue(SpanEventType.START, span.model_copy(deep=True), self.processors) @@ -276,6 +304,9 @@ async def end_span( if span.end_time is None: span.end_time = datetime.now(UTC) + # Close the dedicated obs wrapper span (detach context + end it). + close_obs_span(self._obs_handles.pop(span.id, None)) + span.input = recursive_model_dump(span.input) if span.input else None span.output = recursive_model_dump(span.output) if span.output else None span.data = recursive_model_dump(span.data) if span.data else None diff --git a/tests/lib/core/tracing/test_obs_ids.py b/tests/lib/core/tracing/test_obs_ids.py new file mode 100644 index 000000000..ddd079743 --- /dev/null +++ b/tests/lib/core/tracing/test_obs_ids.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +import sys +import types + +import pytest + +from agentex.lib.core.tracing import obs_ids +from agentex.lib.core.tracing.obs_ids import get_obs_mode, obs_correlation + + +class TestGetObsMode: + @pytest.mark.parametrize( + "raw, expected", + [ + (None, "dd_only"), # unset + ("", "dd_only"), # empty + ("dd_only", "dd_only"), + ("lgtm", "lgtm"), + ("LGTM", "lgtm"), # case-insensitive + (" lgtm ", "lgtm"), # trimmed + ("dual", "dd_only"), # removed mode -> safe degrade + ("garbage", "dd_only"), # unrecognized -> safe degrade + ], + ) + def test_mode_resolution(self, monkeypatch, raw, expected): + if raw is None: + monkeypatch.delenv("SGP_OBS_MODE", raising=False) + else: + monkeypatch.setenv("SGP_OBS_MODE", raw) + assert get_obs_mode() == expected + + +class TestObsCorrelation: + def test_lgtm_mode_reads_otel_and_emits_underscored_keys(self, monkeypatch): + monkeypatch.setenv("SGP_OBS_MODE", "lgtm") + monkeypatch.setattr(obs_ids, "_lgtm_ids", lambda: ("otel_trace", "otel_span")) + # In lgtm mode ddtrace must NOT be consulted. + monkeypatch.setattr( + obs_ids, "_ddtrace_ids", lambda: pytest.fail("ddtrace read in lgtm mode") + ) + + assert obs_correlation() == { + "obs_trace_id": "otel_trace", + "obs_span_id": "otel_span", + } + + def test_dd_only_mode_reads_ddtrace(self, monkeypatch): + monkeypatch.setenv("SGP_OBS_MODE", "dd_only") + monkeypatch.setattr(obs_ids, "_ddtrace_ids", lambda: ("dd_trace", "dd_span")) + monkeypatch.setattr( + obs_ids, "_lgtm_ids", lambda: pytest.fail("otel read in dd_only mode") + ) + + assert obs_correlation() == { + "obs_trace_id": "dd_trace", + "obs_span_id": "dd_span", + } + + def test_stale_dual_degrades_to_ddtrace(self, monkeypatch): + """A leftover SGP_OBS_MODE=dual must behave as dd_only, not read OTel.""" + monkeypatch.setenv("SGP_OBS_MODE", "dual") + monkeypatch.setattr(obs_ids, "_ddtrace_ids", lambda: ("dd_trace", "dd_span")) + monkeypatch.setattr( + obs_ids, "_lgtm_ids", lambda: pytest.fail("otel read for stale dual mode") + ) + + assert obs_correlation() == { + "obs_trace_id": "dd_trace", + "obs_span_id": "dd_span", + } + + def test_no_active_context_returns_empty(self, monkeypatch): + monkeypatch.delenv("SGP_OBS_MODE", raising=False) # dd_only + monkeypatch.setattr(obs_ids, "_ddtrace_ids", lambda: None) + + assert obs_correlation() == {} + + def test_resolver_exception_is_swallowed(self, monkeypatch): + """A misbehaving tracer must not propagate out of obs_correlation.""" + monkeypatch.delenv("SGP_OBS_MODE", raising=False) # dd_only + + def boom(): + raise RuntimeError("tracer blew up") + + monkeypatch.setattr(obs_ids, "_ddtrace_ids", boom) + assert obs_correlation() == {} + + +class TestIdFormatting: + """Pin the W3C hex shape (32-hex trace, 16-hex span) of the resolvers.""" + + def test_ddtrace_ids_formats_w3c_hex(self, monkeypatch): + ctx = types.SimpleNamespace(trace_id=0xABC, span_id=0xFF) + tracer = types.SimpleNamespace(current_trace_context=lambda: ctx) + fake_ddtrace = types.ModuleType("ddtrace") + fake_trace = types.ModuleType("ddtrace.trace") + fake_trace.tracer = tracer + monkeypatch.setitem(sys.modules, "ddtrace", fake_ddtrace) + monkeypatch.setitem(sys.modules, "ddtrace.trace", fake_trace) + + trace_id, span_id = obs_ids._ddtrace_ids() + assert trace_id == "00000000000000000000000000000abc" + assert span_id == "000000000000000000ff"[-16:] # 16-hex + assert len(trace_id) == 32 and len(span_id) == 16 + + def test_lgtm_ids_formats_w3c_hex(self, monkeypatch): + span_ctx = types.SimpleNamespace(trace_id=0xABC, span_id=0xFF, is_valid=True) + current_span = types.SimpleNamespace(get_span_context=lambda: span_ctx) + fake_trace_mod = types.SimpleNamespace(get_current_span=lambda: current_span) + fake_otel = types.ModuleType("opentelemetry") + fake_otel.trace = fake_trace_mod + monkeypatch.setitem(sys.modules, "opentelemetry", fake_otel) + + trace_id, span_id = obs_ids._lgtm_ids() + assert trace_id == "00000000000000000000000000000abc" + assert len(trace_id) == 32 and len(span_id) == 16 + + def test_ddtrace_ids_none_when_no_context(self, monkeypatch): + tracer = types.SimpleNamespace(current_trace_context=lambda: None) + fake_ddtrace = types.ModuleType("ddtrace") + fake_trace = types.ModuleType("ddtrace.trace") + fake_trace.tracer = tracer + monkeypatch.setitem(sys.modules, "ddtrace", fake_ddtrace) + monkeypatch.setitem(sys.modules, "ddtrace.trace", fake_trace) + + assert obs_ids._ddtrace_ids() is None diff --git a/tests/lib/core/tracing/test_obs_span.py b/tests/lib/core/tracing/test_obs_span.py new file mode 100644 index 000000000..2638156a2 --- /dev/null +++ b/tests/lib/core/tracing/test_obs_span.py @@ -0,0 +1,399 @@ +from __future__ import annotations + +import sys +import types +from unittest.mock import MagicMock + +from agentex.lib.core.tracing import obs_span +from agentex.lib.core.tracing.trace import Trace + + +# --------------------------------------------------------------------------- # +# Fake OTel (lgtm) and fake ddtrace (dd_only) SDKs injected via sys.modules. +# --------------------------------------------------------------------------- # +class _FakeSpanContext: + def __init__(self, trace_id: int, span_id: int, is_valid: bool = True): + self.trace_id = trace_id + self.span_id = span_id + self.is_valid = is_valid + + +class _FakeOtelSpan: + def __init__(self, name: str, trace_id: int, span_id: int): + self.name = name + self._ctx = _FakeSpanContext(trace_id, span_id) + self.ended = False + self.attributes: dict = {} + + def set_attribute(self, key, value): + self.attributes[key] = value + + def get_span_context(self): + return self._ctx + + def end(self): + self.ended = True + + +def _install_fake_otel(monkeypatch, *, trace_id=0xABC, span_id=0xFF): + record: dict = {"span": None, "attached": [], "detached": []} + + def start_span(name): + span = _FakeOtelSpan(name, trace_id, span_id) + record["span"] = span + return span + + tracer = types.SimpleNamespace(start_span=start_span) + fake_trace = types.SimpleNamespace( + get_tracer=lambda _name: tracer, + set_span_in_context=lambda span: {"span": span}, + ) + fake_context = types.SimpleNamespace( + attach=lambda ctx: record["attached"].append(ctx) or object(), + detach=lambda token: record["detached"].append(token), + ) + fake_otel = types.ModuleType("opentelemetry") + fake_otel.trace = fake_trace + fake_otel.context = fake_context + monkeypatch.setitem(sys.modules, "opentelemetry", fake_otel) + return record + + +class _FakeDDSpan: + def __init__(self, name: str, trace_id: int, span_id: int): + self.name = name + self.trace_id = trace_id + self.span_id = span_id + self.finished = False + self.tags: dict = {} + + def set_tag(self, key, value): + self.tags[key] = value + + def finish(self): + self.finished = True + + +def _install_fake_ddtrace(monkeypatch, *, active=True, trace_id=0xABC, span_id=0xFF): + record: dict = {"span": None, "started": []} + + def start_span(name, activate=False): + span = _FakeDDSpan(name, trace_id, span_id) + record["span"] = span + record["started"].append((name, activate)) + return span + + tracer = types.SimpleNamespace( + current_trace_context=lambda: (object() if active else None), + start_span=start_span, + ) + fake_ddtrace = types.ModuleType("ddtrace") + fake_trace = types.ModuleType("ddtrace.trace") + fake_trace.tracer = tracer + monkeypatch.setitem(sys.modules, "ddtrace", fake_ddtrace) + monkeypatch.setitem(sys.modules, "ddtrace.trace", fake_trace) + return record + + +# --------------------------------------------------------------------------- # +# lgtm -> OTel wrapper +# --------------------------------------------------------------------------- # +class TestOtelWrapper: + def test_lgtm_opens_named_span_and_reads_its_ids(self, monkeypatch): + monkeypatch.setenv("SGP_OBS_MODE", "lgtm") + record = _install_fake_otel(monkeypatch, trace_id=0xABC, span_id=0xFF) + + handle = obs_span.open_obs_span( + "rocket.tool.fetch", business_span_id="bspan-1", business_trace_id="btrace-1" + ) + + assert handle is not None + assert record["span"].name == "rocket.tool.fetch" # named for the step + assert len(record["attached"]) == 1 # made active + assert handle.correlation == { + "obs_trace_id": "00000000000000000000000000000abc", + "obs_span_id": "000000000000000000ff"[-16:], + } + # reverse tag: business ids stamped on the obs span + assert record["span"].attributes == { + "agentex.business_span_id": "bspan-1", + "agentex.business_trace_id": "btrace-1", + } + + def test_invalid_span_context_yields_empty_correlation(self, monkeypatch): + monkeypatch.setenv("SGP_OBS_MODE", "lgtm") + _install_fake_otel(monkeypatch) + + def start_span(name): + span = _FakeOtelSpan(name, 0, 0) + span._ctx = _FakeSpanContext(0, 0, is_valid=False) + return span + + sys.modules["opentelemetry"].trace.get_tracer = lambda _n: types.SimpleNamespace( + start_span=start_span + ) + handle = obs_span.open_obs_span("step") + assert handle is not None + assert handle.correlation == {} + + def test_close_detaches_and_ends(self, monkeypatch): + monkeypatch.setenv("SGP_OBS_MODE", "lgtm") + record = _install_fake_otel(monkeypatch) + handle = obs_span.open_obs_span("step") + + obs_span.close_obs_span(handle) + + assert record["span"].ended is True + assert len(record["detached"]) == 1 + + def test_close_none_is_noop(self): + obs_span.close_obs_span(None) # must not raise + + +# --------------------------------------------------------------------------- # +# dd_only -> ddtrace wrapper (only when a request trace is active) +# --------------------------------------------------------------------------- # +class TestDdtraceWrapper: + def test_dd_only_with_active_ctx_opens_named_span(self, monkeypatch): + monkeypatch.setenv("SGP_OBS_MODE", "dd_only") + record = _install_fake_ddtrace(monkeypatch, active=True, trace_id=0xABC, span_id=0xFF) + + handle = obs_span.open_obs_span( + "rocket.tool.fetch", business_span_id="bspan-9", business_trace_id="btrace-9" + ) + + assert handle is not None + assert record["span"].name == "rocket.tool.fetch" + assert record["started"] == [("rocket.tool.fetch", True)] # activated + assert handle.correlation == { + "obs_trace_id": "00000000000000000000000000000abc", + "obs_span_id": "000000000000000000ff"[-16:], + } + # reverse tag on the ddtrace span + assert record["span"].tags == { + "agentex.business_span_id": "bspan-9", + "agentex.business_trace_id": "btrace-9", + } + + obs_span.close_obs_span(handle) + assert record["span"].finished is True + + def test_dd_only_without_active_ctx_returns_none(self, monkeypatch): + """Bare-uvicorn / no ddtrace-run: nothing active -> no orphan wrapper.""" + monkeypatch.setenv("SGP_OBS_MODE", "dd_only") + record = _install_fake_ddtrace(monkeypatch, active=False) + + assert obs_span.open_obs_span("step") is None + assert record["span"] is None # never created a span + + +# --------------------------------------------------------------------------- # +# End-to-end through Trace.start_span / end_span +# --------------------------------------------------------------------------- # +class TestTraceIntegration: + def test_lgtm_business_span_tagged_with_wrapper_ids(self, monkeypatch): + monkeypatch.setenv("SGP_OBS_MODE", "lgtm") + record = _install_fake_otel(monkeypatch, trace_id=0x111, span_id=0x222) + + trace = Trace(processors=[], client=MagicMock(), trace_id="task-run-1") + span = trace.start_span(name="chat_completion") + + assert record["span"].name == "chat_completion" # dedicated named span + assert span.data["obs_trace_id"] == "00000000000000000000000000000111" + assert span.data["obs_span_id"] == "0000000000000222" + assert span.trace_id == "task-run-1" # business id unchanged + assert span.id in trace._obs_handles + # bidirectional: the obs span carries the business ids (reverse tag), + # and the business span carries the obs ids (forward edge). + assert record["span"].attributes == { + "agentex.business_span_id": span.id, + "agentex.business_trace_id": "task-run-1", + } + + trace.end_span(span) + assert record["span"].ended is True + assert span.id not in trace._obs_handles + + def test_dd_only_business_span_tagged_via_ddtrace(self, monkeypatch): + monkeypatch.setenv("SGP_OBS_MODE", "dd_only") + record = _install_fake_ddtrace(monkeypatch, active=True, trace_id=0x111, span_id=0x222) + + trace = Trace(processors=[], client=MagicMock(), trace_id="task-run-2") + span = trace.start_span(name="get_state") + + assert record["span"].name == "get_state" + assert span.data["obs_trace_id"] == "00000000000000000000000000000111" + assert span.data["obs_span_id"] == "0000000000000222" + assert record["span"].tags == { + "agentex.business_span_id": span.id, + "agentex.business_trace_id": "task-run-2", + } + + trace.end_span(span) + assert record["span"].finished is True + assert span.id not in trace._obs_handles + + 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: {}) + + trace = Trace(processors=[], client=MagicMock(), trace_id="task-run-3") + span = trace.start_span(name="get_state") + + assert trace._obs_handles == {} # no wrapper opened + assert span.data is None # nothing tagged + trace.end_span(span) # must not raise + + +# --------------------------------------------------------------------------- # +# Non-interference: the two backends are mutually exclusive per mode. +# --------------------------------------------------------------------------- # +class TestNonInterference: + def test_lgtm_touches_only_otel(self, monkeypatch): + monkeypatch.setenv("SGP_OBS_MODE", "lgtm") + otel = _install_fake_otel(monkeypatch) + dd = _install_fake_ddtrace(monkeypatch, active=True) + + obs_span.open_obs_span("step") + + assert otel["span"] is not None # OTel wrapper opened + assert dd["span"] is None # ddtrace never touched + + def test_dd_only_touches_only_ddtrace(self, monkeypatch): + monkeypatch.setenv("SGP_OBS_MODE", "dd_only") + otel = _install_fake_otel(monkeypatch) + dd = _install_fake_ddtrace(monkeypatch, active=True) + + obs_span.open_obs_span("step") + + assert dd["span"] is not None # ddtrace wrapper opened + assert otel["span"] is None # OTel never touched + + +# --------------------------------------------------------------------------- # +# No-op when unconfigured, and never fails the app call. +# --------------------------------------------------------------------------- # +class TestNeverFails: + def test_lgtm_no_otel_installed_returns_none(self, monkeypatch): + monkeypatch.setenv("SGP_OBS_MODE", "lgtm") + monkeypatch.setitem(sys.modules, "opentelemetry", None) # import -> ImportError + assert obs_span.open_obs_span("step") is None + + def test_dd_only_no_ddtrace_installed_returns_none(self, monkeypatch): + monkeypatch.setenv("SGP_OBS_MODE", "dd_only") + monkeypatch.setitem(sys.modules, "ddtrace.trace", None) # import -> ImportError + assert obs_span.open_obs_span("step") is None + + def test_backend_exception_is_swallowed(self, monkeypatch): + monkeypatch.setenv("SGP_OBS_MODE", "lgtm") + _install_fake_otel(monkeypatch) + + def boom(_name): + raise RuntimeError("tracer blew up") + + sys.modules["opentelemetry"].trace.get_tracer = boom + assert obs_span.open_obs_span("step") is None # inner guard + + def test_top_level_guard_swallows_get_mode_error(self, monkeypatch): + # Even if mode resolution itself raises, open_obs_span must not. + monkeypatch.setattr(obs_span, "get_obs_mode", lambda: (_ for _ in ()).throw(RuntimeError())) + assert obs_span.open_obs_span("step") is None + + def test_close_swallows_closer_error(self): + handle = obs_span.ObsSpanHandle({}, lambda: (_ for _ in ()).throw(RuntimeError())) + obs_span.close_obs_span(handle) # must not raise + + def test_unconfigured_lgtm_yields_usable_span_no_raise(self, monkeypatch): + # lgtm requested but OTel not installed: the REAL open_obs_span returns + # None, obs_correlation() returns {} (also no tracer) -> the business + # span is created and fully usable, and nothing raised. + monkeypatch.setenv("SGP_OBS_MODE", "lgtm") + monkeypatch.setitem(sys.modules, "opentelemetry", None) + + trace = Trace(processors=[], client=MagicMock(), trace_id="task-run-4") + span = trace.start_span(name="safe") + + assert span.trace_id == "task-run-4" + assert trace._obs_handles == {} # no wrapper + trace.end_span(span) # must not raise + + +def _install_fake_otel_sequence(monkeypatch, *, trace_id: int, first_span_id: int): + """Fake OTel whose wrapper spans all share ``trace_id`` (children of the one + turn/request obs trace) but get sequential distinct span ids.""" + state: dict = {"next": first_span_id, "spans": []} + + def start_span(name): + sid = state["next"] + state["next"] += 1 + span = _FakeOtelSpan(name, trace_id, sid) + state["spans"].append(span) + return span + + tracer = types.SimpleNamespace(start_span=start_span) + fake_trace = types.SimpleNamespace( + get_tracer=lambda _name: tracer, + set_span_in_context=lambda span: {"span": span}, + ) + fake_context = types.SimpleNamespace( + attach=lambda ctx: object(), + detach=lambda token: None, + ) + fake_otel = types.ModuleType("opentelemetry") + fake_otel.trace = fake_trace + fake_otel.context = fake_context + monkeypatch.setitem(sys.modules, "opentelemetry", fake_otel) + return state + + +class TestTurn2Example: + """Maps the 3-turn mortgage example, Turn 2 (obs trace B): + + get_state -> wrapper wB1 -> obs_span_id = wB1 + retrieve_docs -> wrapper wB2 -> obs_span_id = wB2 + chat_completion -> wrapper wB3 -> obs_span_id = wB3 + create_message -> wrapper wB4 -> obs_span_id = wB4 + + Each step opens its OWN dedicated span named for the step; all four share the + one turn obs trace B, but obs_span_id is distinct per step (not all rB). + """ + + def test_turn2_each_step_gets_distinct_named_wrapper_under_trace_B(self, monkeypatch): + monkeypatch.setenv("SGP_OBS_MODE", "lgtm") + # Turn 2's request obs trace = B (0xB); wrappers get span ids 0xB1.. . + state = _install_fake_otel_sequence(monkeypatch, trace_id=0xB, first_span_id=0xB1) + + run_id = "task-run-mortgage" # business trace_id = the run/task id + trace = Trace(processors=[], client=MagicMock(), trace_id=run_id) + + steps = ["get_state", "retrieve_docs", "chat_completion", "create_message"] + business = [] + for step in steps: + with trace.span(name=step) as s: + business.append(s) + + obs_trace_B = format(0xB, "032x") + expected_obs_span = [format(sid, "016x") for sid in (0xB1, 0xB2, 0xB3, 0xB4)] + + # one dedicated wrapper per step, named for the step, in order + assert [w.name for w in state["spans"]] == steps + + for biz, wrapper, exp_span in zip(business, state["spans"], expected_obs_span): + # forward edge: business span carries the wrapper's ids + assert biz.data["obs_trace_id"] == obs_trace_B # all under trace B + assert biz.data["obs_span_id"] == exp_span # distinct wBn + # reverse tag: wrapper carries the business ids + assert wrapper.attributes == { + "agentex.business_span_id": biz.id, + "agentex.business_trace_id": run_id, + } + + # the whole point of the fix: obs_span_id is DISTINCT per step ... + obs_span_ids = [b.data["obs_span_id"] for b in business] + assert obs_span_ids == expected_obs_span + assert len(set(obs_span_ids)) == 4 + # ... while all four share the single turn obs trace B + assert {b.data["obs_trace_id"] for b in business} == {obs_trace_B} + # business trace stays the run/task id, not the obs trace + assert {b.trace_id for b in business} == {run_id} From c6106394548902e565625c5404d5ea1c95a4635f Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Sat, 1 Aug 2026 12:29:32 -0700 Subject: [PATCH 02/12] observability: propagate business-span error status to wrapper obs span Observability review (P1): the dedicated wrapper obs span was ended/finished without recording failure, so a failed business step (e.g. chat_completion) showed green in Tempo/DD -- violating "observe both success and failure" and undercutting the meaningful-obs_span_id goal. close_obs_span now takes the business span's error (from get_span_error) and marks the obs span before closing: - OTel: span.set_status(Status(ERROR, msg)) + error.type attribute - ddtrace: span.error = 1 + error.type / error.message tags end_span passes error=get_span_error(span) on both sync and async paths. Success path is unchanged (no status set). Guarded so error-marking can never break the close. Co-Authored-By: Claude Opus 4.8 --- src/agentex/lib/core/tracing/obs_span.py | 48 ++++++++++++++--- src/agentex/lib/core/tracing/trace.py | 12 +++-- tests/lib/core/tracing/test_obs_span.py | 69 ++++++++++++++++++++++++ 3 files changed, 116 insertions(+), 13 deletions(-) diff --git a/src/agentex/lib/core/tracing/obs_span.py b/src/agentex/lib/core/tracing/obs_span.py index 77186d2ad..4d0b6c1db 100644 --- a/src/agentex/lib/core/tracing/obs_span.py +++ b/src/agentex/lib/core/tracing/obs_span.py @@ -47,7 +47,11 @@ class ObsSpanHandle: __slots__ = ("correlation", "_close") - def __init__(self, correlation: Dict[str, str], close: Callable[[], None]): + def __init__( + self, + correlation: Dict[str, str], + close: Callable[[Optional[Dict[str, str]]], None], + ): self.correlation = correlation self._close = close @@ -79,11 +83,21 @@ def _open_otel_span( sc = span.get_span_context() correlation = _hex_ids(sc.trace_id, sc.span_id) if (sc and sc.is_valid) else {} - def _close() -> None: + def _close(error: Optional[Dict[str, str]] = None) -> None: try: - context.detach(token) + if error: + # Reflect the business-step failure on the obs span so it + # isn't a false green when you pivot from a failed span. + span.set_status( + trace.Status(trace.StatusCode.ERROR, error.get("message")) + ) + if error.get("type"): + span.set_attribute("error.type", error["type"]) finally: - span.end() + try: + context.detach(token) + finally: + span.end() return ObsSpanHandle(correlation, _close) except Exception: # pragma: no cover - best-effort; never break the business span @@ -110,7 +124,20 @@ def _open_ddtrace_span( if business_trace_id: span.set_tag(_ATTR_BUSINESS_TRACE_ID, business_trace_id) correlation = _hex_ids(span.trace_id, span.span_id) if span.trace_id else {} - return ObsSpanHandle(correlation, span.finish) + + def _close(error: Optional[Dict[str, str]] = None) -> None: + try: + if error: + # Reflect the business-step failure on the obs span. + span.error = 1 + if error.get("type"): + span.set_tag("error.type", error["type"]) + if error.get("message"): + span.set_tag("error.message", error["message"]) + finally: + span.finish() + + return ObsSpanHandle(correlation, _close) except Exception: # pragma: no cover - best-effort return None @@ -143,11 +170,16 @@ def open_obs_span( return None -def close_obs_span(handle: Optional[ObsSpanHandle]) -> None: - """Close the wrapper span (detach + end, or finish). Safe on ``None``.""" +def close_obs_span( + handle: Optional[ObsSpanHandle], + error: Optional[Dict[str, str]] = None, +) -> None: + """Close the wrapper span (detach + end, or finish). When ``error`` is given + (the business span failed), mark the obs span errored first so it reflects + failure rather than a false green. Safe on ``None``.""" if handle is None: return try: - handle._close() + handle._close(error) except Exception: # pragma: no cover - best-effort pass diff --git a/src/agentex/lib/core/tracing/trace.py b/src/agentex/lib/core/tracing/trace.py index 031e69e9c..7ac6aa2ef 100644 --- a/src/agentex/lib/core/tracing/trace.py +++ b/src/agentex/lib/core/tracing/trace.py @@ -17,7 +17,7 @@ close_obs_span, open_obs_span, ) -from agentex.lib.core.tracing.span_error import set_span_error +from agentex.lib.core.tracing.span_error import get_span_error, set_span_error from agentex.lib.core.tracing.span_queue import ( SpanEventType, AsyncSpanQueue, @@ -135,8 +135,9 @@ def end_span( if span.end_time is None: span.end_time = datetime.now(UTC) - # Close the dedicated obs wrapper span (detach context + end it). - close_obs_span(self._obs_handles.pop(span.id, None)) + # Close the dedicated obs wrapper span; propagate the business-span error + # (if any) so the obs span reflects failure, not a false green. + close_obs_span(self._obs_handles.pop(span.id, None), error=get_span_error(span)) span.input = recursive_model_dump(span.input) if span.input else None span.output = recursive_model_dump(span.output) if span.output else None @@ -304,8 +305,9 @@ async def end_span( if span.end_time is None: span.end_time = datetime.now(UTC) - # Close the dedicated obs wrapper span (detach context + end it). - close_obs_span(self._obs_handles.pop(span.id, None)) + # Close the dedicated obs wrapper span; propagate the business-span error + # (if any) so the obs span reflects failure, not a false green. + close_obs_span(self._obs_handles.pop(span.id, None), error=get_span_error(span)) span.input = recursive_model_dump(span.input) if span.input else None span.output = recursive_model_dump(span.output) if span.output else None diff --git a/tests/lib/core/tracing/test_obs_span.py b/tests/lib/core/tracing/test_obs_span.py index 2638156a2..aed5e2abc 100644 --- a/tests/lib/core/tracing/test_obs_span.py +++ b/tests/lib/core/tracing/test_obs_span.py @@ -4,6 +4,8 @@ import types from unittest.mock import MagicMock +import pytest + from agentex.lib.core.tracing import obs_span from agentex.lib.core.tracing.trace import Trace @@ -18,16 +20,30 @@ def __init__(self, trace_id: int, span_id: int, is_valid: bool = True): self.is_valid = is_valid +class _FakeStatusCode: + ERROR = "ERROR" + OK = "OK" + UNSET = "UNSET" + + +def _FakeStatus(code, description=None): + return {"code": code, "description": description} + + class _FakeOtelSpan: def __init__(self, name: str, trace_id: int, span_id: int): self.name = name self._ctx = _FakeSpanContext(trace_id, span_id) self.ended = False self.attributes: dict = {} + self.status = None def set_attribute(self, key, value): self.attributes[key] = value + def set_status(self, status): + self.status = status + def get_span_context(self): return self._ctx @@ -47,6 +63,8 @@ def start_span(name): fake_trace = types.SimpleNamespace( get_tracer=lambda _name: tracer, set_span_in_context=lambda span: {"span": span}, + Status=_FakeStatus, + StatusCode=_FakeStatusCode, ) fake_context = types.SimpleNamespace( attach=lambda ctx: record["attached"].append(ctx) or object(), @@ -65,6 +83,7 @@ def __init__(self, name: str, trace_id: int, span_id: int): self.trace_id = trace_id self.span_id = span_id self.finished = False + self.error = 0 self.tags: dict = {} def set_tag(self, key, value): @@ -149,6 +168,27 @@ def test_close_detaches_and_ends(self, monkeypatch): def test_close_none_is_noop(self): obs_span.close_obs_span(None) # must not raise + def test_close_with_error_marks_otel_status(self, monkeypatch): + monkeypatch.setenv("SGP_OBS_MODE", "lgtm") + record = _install_fake_otel(monkeypatch) + handle = obs_span.open_obs_span("step") + + obs_span.close_obs_span(handle, error={"type": "ValueError", "message": "boom"}) + + assert record["span"].status == {"code": "ERROR", "description": "boom"} + assert record["span"].attributes.get("error.type") == "ValueError" + assert record["span"].ended is True + + def test_close_without_error_leaves_otel_status_unset(self, monkeypatch): + monkeypatch.setenv("SGP_OBS_MODE", "lgtm") + record = _install_fake_otel(monkeypatch) + handle = obs_span.open_obs_span("step") + + obs_span.close_obs_span(handle) # success path + + assert record["span"].status is None + assert record["span"].ended is True + # --------------------------------------------------------------------------- # # dd_only -> ddtrace wrapper (only when a request trace is active) @@ -186,6 +226,18 @@ def test_dd_only_without_active_ctx_returns_none(self, monkeypatch): assert obs_span.open_obs_span("step") is None assert record["span"] is None # never created a span + def test_close_with_error_marks_ddtrace_span(self, monkeypatch): + monkeypatch.setenv("SGP_OBS_MODE", "dd_only") + record = _install_fake_ddtrace(monkeypatch, active=True) + handle = obs_span.open_obs_span("step") + + obs_span.close_obs_span(handle, error={"type": "ValueError", "message": "boom"}) + + assert record["span"].error == 1 + assert record["span"].tags.get("error.type") == "ValueError" + assert record["span"].tags.get("error.message") == "boom" + assert record["span"].finished is True + # --------------------------------------------------------------------------- # # End-to-end through Trace.start_span / end_span @@ -233,6 +285,21 @@ def test_dd_only_business_span_tagged_via_ddtrace(self, monkeypatch): assert record["span"].finished is True assert span.id not in trace._obs_handles + def test_lgtm_wrapper_marked_error_when_business_step_raises(self, monkeypatch): + monkeypatch.setenv("SGP_OBS_MODE", "lgtm") + record = _install_fake_otel(monkeypatch, trace_id=0x111, span_id=0x222) + + trace = Trace(processors=[], client=MagicMock(), trace_id="task-run-err") + with pytest.raises(ValueError): + with trace.span(name="chat_completion"): + raise ValueError("boom") + + # the failed step's obs span reflects the failure, not a false green + assert record["span"].name == "chat_completion" + assert record["span"].ended is True + assert record["span"].status == {"code": "ERROR", "description": "boom"} + assert record["span"].attributes.get("error.type") == "ValueError" + 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) @@ -335,6 +402,8 @@ def start_span(name): fake_trace = types.SimpleNamespace( get_tracer=lambda _name: tracer, set_span_in_context=lambda span: {"span": span}, + Status=_FakeStatus, + StatusCode=_FakeStatusCode, ) fake_context = types.SimpleNamespace( attach=lambda ctx: object(), From 3edf16d725cfbea9804737d966d6f98cd7cde440 Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Sat, 1 Aug 2026 22:36:59 -0700 Subject: [PATCH 03/12] fix(tracing): nest ddtrace wrapper under the active context (child_of) ddtrace start_span does not auto-parent (unlike OTel): start_span(name) mints a new ROOT trace every call, so a turn's business spans scattered across N Datadog traces (verified live: 52 spans -> 52 distinct obs_trace_ids). Pass child_of=current_trace_context() so wrappers nest under the request/turn trace and roll up into one trace; obs_span_id stays distinct per step. Co-Authored-By: Claude Opus 4.8 --- src/agentex/lib/core/tracing/obs_span.py | 10 ++++++++-- tests/lib/core/tracing/test_obs_span.py | 15 +++++++++++---- 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/src/agentex/lib/core/tracing/obs_span.py b/src/agentex/lib/core/tracing/obs_span.py index 4d0b6c1db..39234b8ff 100644 --- a/src/agentex/lib/core/tracing/obs_span.py +++ b/src/agentex/lib/core/tracing/obs_span.py @@ -116,9 +116,15 @@ def _open_ddtrace_span( try: # Only wrap when ddtrace is actually tracing the request; otherwise a # wrapper would be an orphan root trace in an un-instrumented process. - if tracer.current_trace_context() is None: + ctx = tracer.current_trace_context() + if ctx is None: return None - span = tracer.start_span(name, activate=True) + # child_of=ctx is load-bearing: ddtrace's start_span does NOT auto-parent + # to the active span (unlike OTel), so start_span(name) alone mints a NEW + # root trace every call -- scattering a turn's business spans across N + # Datadog traces. Parenting to the active request/turn context rolls them + # into one trace while obs_span_id stays distinct per step. + span = tracer.start_span(name, child_of=ctx, activate=True) if business_span_id: span.set_tag(_ATTR_BUSINESS_SPAN_ID, business_span_id) if business_trace_id: diff --git a/tests/lib/core/tracing/test_obs_span.py b/tests/lib/core/tracing/test_obs_span.py index aed5e2abc..c35fbe1a0 100644 --- a/tests/lib/core/tracing/test_obs_span.py +++ b/tests/lib/core/tracing/test_obs_span.py @@ -95,15 +95,17 @@ def finish(self): def _install_fake_ddtrace(monkeypatch, *, active=True, trace_id=0xABC, span_id=0xFF): record: dict = {"span": None, "started": []} + ctx_obj = object() if active else None + record["ctx"] = ctx_obj - def start_span(name, activate=False): + def start_span(name, child_of=None, activate=False): span = _FakeDDSpan(name, trace_id, span_id) record["span"] = span - record["started"].append((name, activate)) + record["started"].append({"name": name, "child_of": child_of, "activate": activate}) return span tracer = types.SimpleNamespace( - current_trace_context=lambda: (object() if active else None), + current_trace_context=lambda: ctx_obj, start_span=start_span, ) fake_ddtrace = types.ModuleType("ddtrace") @@ -204,7 +206,12 @@ def test_dd_only_with_active_ctx_opens_named_span(self, monkeypatch): assert handle is not None assert record["span"].name == "rocket.tool.fetch" - assert record["started"] == [("rocket.tool.fetch", True)] # activated + started = record["started"][0] + assert started["name"] == "rocket.tool.fetch" + assert started["activate"] is True + # child_of is the active request/turn context -> the wrapper nests under + # it instead of minting a new root trace (ddtrace does not auto-parent). + assert started["child_of"] is record["ctx"] assert handle.correlation == { "obs_trace_id": "00000000000000000000000000000abc", "obs_span_id": "000000000000000000ff"[-16:], From f806fc7c30c0ade38e61bbc0e2124ff514c91e74 Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Sun, 2 Aug 2026 16:01:48 -0700 Subject: [PATCH 04/12] fix(tracing): track obs wrapper spans in a module-level registry The obs wrapper span (opened in start_span, ended in end_span) was tracked in an INSTANCE dict (self._obs_handles). But TracingService creates a fresh Trace object for every call -- self._tracer.trace(trace_id) in BOTH start_span and end_span -- so end_span ran on a different instance with an empty dict: the handle was never found, close_obs_span(None) was a no-op, and the OTel/ddtrace wrapper span was never .end()ed. Consequence in lgtm mode: the wrapper span records and its ids are written to Postgres (read at start), but since Simple/Batch span processors only export on span end, the span never reaches Tempo -- the turn trace was silently missing while everything looked correct (provider ours, sampler ALWAYS_ON, recording=True, ids stored). Fix: move the handle registry to module level, keyed by the uuid4 span id, so it survives across Trace instances. Adds a regression test that starts a span on one Trace instance and ends it on another. Co-Authored-By: Claude Opus 4.8 --- src/agentex/lib/core/tracing/trace.py | 29 +++++++++++----- tests/lib/core/tracing/test_obs_span.py | 45 ++++++++++++++++++++++--- 2 files changed, 61 insertions(+), 13 deletions(-) diff --git a/src/agentex/lib/core/tracing/trace.py b/src/agentex/lib/core/tracing/trace.py index 7ac6aa2ef..b8c3a6fc0 100644 --- a/src/agentex/lib/core/tracing/trace.py +++ b/src/agentex/lib/core/tracing/trace.py @@ -30,6 +30,17 @@ logger = make_logger(__name__) +# Live per-business-span obs wrapper spans, keyed by the (uuid4) business span id, +# in a MODULE-LEVEL registry -- deliberately NOT on the Trace/AsyncTrace instance. +# TracingService creates a FRESH trace object for every call +# (`self._tracer.trace(trace_id)` in both start_span and end_span), so an +# instance-local dict loses the handle between start and end: end_span's new +# instance can't find it, close_obs_span(None) is a no-op, and the OTel wrapper +# span is never .end()ed -> never exported (Simple/Batch processors only emit on +# end). A module-level dict keyed by the unique span id survives across instances; +# uuid4 span ids cannot collide across concurrent traces. +_OBS_HANDLES: dict[str, ObsSpanHandle] = {} + class Trace: """ @@ -54,8 +65,9 @@ def __init__( self.processors = processors self.client = client self.trace_id = trace_id - # Live per-business-span obs wrapper spans, keyed by business span id. - self._obs_handles: dict[str, ObsSpanHandle] = {} + # Obs wrapper spans are tracked in the module-level _OBS_HANDLES registry + # (see comment there): a fresh trace object is created per start/end call, + # so the handle must not live on the instance. def start_span( self, @@ -112,7 +124,7 @@ def start_span( task_id=task_id, ) if obs_handle is not None: - self._obs_handles[span.id] = obs_handle + _OBS_HANDLES[span.id] = obs_handle for processor in self.processors: processor.on_span_start(span) @@ -137,7 +149,7 @@ def end_span( # Close the dedicated obs wrapper span; propagate the business-span error # (if any) so the obs span reflects failure, not a false green. - close_obs_span(self._obs_handles.pop(span.id, None), error=get_span_error(span)) + close_obs_span(_OBS_HANDLES.pop(span.id, None), error=get_span_error(span)) span.input = recursive_model_dump(span.input) if span.input else None span.output = recursive_model_dump(span.output) if span.output else None @@ -225,8 +237,9 @@ def __init__( self.client = client self.trace_id = trace_id self._span_queue = span_queue or get_default_span_queue() - # Live per-business-span obs wrapper spans, keyed by business span id. - self._obs_handles: dict[str, ObsSpanHandle] = {} + # Obs wrapper spans are tracked in the module-level _OBS_HANDLES registry + # (see comment there): a fresh trace object is created per start/end call, + # so the handle must not live on the instance. async def start_span( self, @@ -282,7 +295,7 @@ async def start_span( task_id=task_id, ) if obs_handle is not None: - self._obs_handles[span.id] = obs_handle + _OBS_HANDLES[span.id] = obs_handle if self.processors: self._span_queue.enqueue(SpanEventType.START, span.model_copy(deep=True), self.processors) @@ -307,7 +320,7 @@ async def end_span( # Close the dedicated obs wrapper span; propagate the business-span error # (if any) so the obs span reflects failure, not a false green. - close_obs_span(self._obs_handles.pop(span.id, None), error=get_span_error(span)) + close_obs_span(_OBS_HANDLES.pop(span.id, None), error=get_span_error(span)) span.input = recursive_model_dump(span.input) if span.input else None span.output = recursive_model_dump(span.output) if span.output else None diff --git a/tests/lib/core/tracing/test_obs_span.py b/tests/lib/core/tracing/test_obs_span.py index c35fbe1a0..aa7d826bb 100644 --- a/tests/lib/core/tracing/test_obs_span.py +++ b/tests/lib/core/tracing/test_obs_span.py @@ -7,9 +7,20 @@ import pytest from agentex.lib.core.tracing import obs_span +from agentex.lib.core.tracing import trace as trace_module from agentex.lib.core.tracing.trace import Trace +@pytest.fixture(autouse=True) +def _clear_obs_handles(): + """The obs-handle registry is module-level (survives across Trace instances, + which is the whole point of the fix). Clear it around each test so leftover + handles never leak between tests.""" + trace_module._OBS_HANDLES.clear() + yield + trace_module._OBS_HANDLES.clear() + + # --------------------------------------------------------------------------- # # Fake OTel (lgtm) and fake ddtrace (dd_only) SDKs injected via sys.modules. # --------------------------------------------------------------------------- # @@ -261,7 +272,7 @@ def test_lgtm_business_span_tagged_with_wrapper_ids(self, monkeypatch): assert span.data["obs_trace_id"] == "00000000000000000000000000000111" assert span.data["obs_span_id"] == "0000000000000222" assert span.trace_id == "task-run-1" # business id unchanged - assert span.id in trace._obs_handles + assert span.id in trace_module._OBS_HANDLES # bidirectional: the obs span carries the business ids (reverse tag), # and the business span carries the obs ids (forward edge). assert record["span"].attributes == { @@ -271,7 +282,31 @@ def test_lgtm_business_span_tagged_with_wrapper_ids(self, monkeypatch): trace.end_span(span) assert record["span"].ended is True - assert span.id not in trace._obs_handles + assert span.id not in trace_module._OBS_HANDLES + + def test_wrapper_ends_across_separate_trace_instances(self, monkeypatch): + # Regression for the export bug: TracingService creates a FRESH trace + # object for start_span AND for end_span (self._tracer.trace(trace_id) in + # both). The obs handle is stored in the module-level registry, so a + # DIFFERENT instance ending the span still finds it and calls .end() on + # the OTel wrapper. With an instance-local dict this regressed: end_span's + # new instance had an empty dict -> close_obs_span(None) -> the wrapper + # span was never ended -> never exported to Tempo (recording, ids stored, + # but absent from the trace backend). + monkeypatch.setenv("SGP_OBS_MODE", "lgtm") + record = _install_fake_otel(monkeypatch, trace_id=0x111, span_id=0x222) + + starter = Trace(processors=[], client=MagicMock(), trace_id="task-run-x") + span = starter.start_span(name="chat_completion") + assert record["span"].ended is False + assert span.id in trace_module._OBS_HANDLES + + # A completely separate Trace instance ends the span. + ender = Trace(processors=[], client=MagicMock(), trace_id="task-run-x") + ender.end_span(span) + + assert record["span"].ended is True # wrapper WAS ended -> exportable + assert span.id not in trace_module._OBS_HANDLES # handle cleaned up def test_dd_only_business_span_tagged_via_ddtrace(self, monkeypatch): monkeypatch.setenv("SGP_OBS_MODE", "dd_only") @@ -290,7 +325,7 @@ def test_dd_only_business_span_tagged_via_ddtrace(self, monkeypatch): trace.end_span(span) assert record["span"].finished is True - assert span.id not in trace._obs_handles + assert span.id not in trace_module._OBS_HANDLES def test_lgtm_wrapper_marked_error_when_business_step_raises(self, monkeypatch): monkeypatch.setenv("SGP_OBS_MODE", "lgtm") @@ -315,7 +350,7 @@ def test_dd_only_no_ctx_falls_back_to_ambient(self, monkeypatch): trace = Trace(processors=[], client=MagicMock(), trace_id="task-run-3") span = trace.start_span(name="get_state") - assert trace._obs_handles == {} # no wrapper opened + assert span.id not in trace_module._OBS_HANDLES # no wrapper opened assert span.data is None # nothing tagged trace.end_span(span) # must not raise @@ -389,7 +424,7 @@ def test_unconfigured_lgtm_yields_usable_span_no_raise(self, monkeypatch): span = trace.start_span(name="safe") assert span.trace_id == "task-run-4" - assert trace._obs_handles == {} # no wrapper + assert span.id not in trace_module._OBS_HANDLES # no wrapper trace.end_span(span) # must not raise From d12184801266f068a2cf3a92142e0b992ec1da2e Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Sun, 2 Aug 2026 22:30:00 -0700 Subject: [PATCH 05/12] style(tracing): ruff format + import-sort obs-correlation files Fixes the failing lint job on this PR (ruff I001 import ordering + format) on the obs_ids / obs_span / trace correlation-edge files and their tests. No logic change. Co-Authored-By: Claude Opus 4.8 --- src/agentex/lib/core/tracing/obs_ids.py | 1 + src/agentex/lib/core/tracing/obs_span.py | 9 ++--- src/agentex/lib/core/tracing/trace.py | 10 ++--- tests/lib/core/tracing/test_obs_ids.py | 24 +++++------- tests/lib/core/tracing/test_obs_span.py | 49 ++++++++++-------------- 5 files changed, 38 insertions(+), 55 deletions(-) diff --git a/src/agentex/lib/core/tracing/obs_ids.py b/src/agentex/lib/core/tracing/obs_ids.py index bac500b20..aa5b2678c 100644 --- a/src/agentex/lib/core/tracing/obs_ids.py +++ b/src/agentex/lib/core/tracing/obs_ids.py @@ -24,6 +24,7 @@ This never fabricates ids -- if no observability context is active, it returns an empty dict and the span is simply not tagged. """ + from __future__ import annotations import os diff --git a/src/agentex/lib/core/tracing/obs_span.py b/src/agentex/lib/core/tracing/obs_span.py index 39234b8ff..37e01b430 100644 --- a/src/agentex/lib/core/tracing/obs_span.py +++ b/src/agentex/lib/core/tracing/obs_span.py @@ -24,9 +24,10 @@ No-op when the relevant tracer isn't importable. Never raises -- observability must never break a business span. """ + from __future__ import annotations -from typing import Callable, Dict, Optional +from typing import Dict, Callable, Optional from agentex.lib.core.tracing.obs_ids import LGTM, get_obs_mode @@ -70,7 +71,7 @@ def _open_otel_span( business_trace_id: Optional[str], ) -> Optional[ObsSpanHandle]: try: - from opentelemetry import context, trace + from opentelemetry import trace, context except ImportError: return None try: @@ -88,9 +89,7 @@ def _close(error: Optional[Dict[str, str]] = None) -> None: if error: # Reflect the business-step failure on the obs span so it # isn't a false green when you pivot from a failed span. - span.set_status( - trace.Status(trace.StatusCode.ERROR, error.get("message")) - ) + span.set_status(trace.Status(trace.StatusCode.ERROR, error.get("message"))) if error.get("type"): span.set_attribute("error.type", error["type"]) finally: diff --git a/src/agentex/lib/core/tracing/trace.py b/src/agentex/lib/core/tracing/trace.py index b8c3a6fc0..0fddc1a98 100644 --- a/src/agentex/lib/core/tracing/trace.py +++ b/src/agentex/lib/core/tracing/trace.py @@ -14,8 +14,8 @@ from agentex.lib.core.tracing.obs_ids import obs_correlation from agentex.lib.core.tracing.obs_span import ( ObsSpanHandle, - close_obs_span, open_obs_span, + close_obs_span, ) from agentex.lib.core.tracing.span_error import get_span_error, set_span_error from agentex.lib.core.tracing.span_queue import ( @@ -106,9 +106,7 @@ def start_span( # obs context (ddtrace) when not in lgtm mode. Business trace_id stays the # run-level task id. id = str(uuid.uuid4()) - obs_handle = open_obs_span( - name, business_span_id=id, business_trace_id=self.trace_id - ) + obs_handle = open_obs_span(name, business_span_id=id, business_trace_id=self.trace_id) obs = obs_handle.correlation if obs_handle is not None else obs_correlation() if obs: serialized_data = {**(serialized_data or {}), **obs} @@ -277,9 +275,7 @@ async def start_span( # obs context (ddtrace) when not in lgtm mode. Business trace_id stays the # run-level task id. id = str(uuid.uuid4()) - obs_handle = open_obs_span( - name, business_span_id=id, business_trace_id=self.trace_id - ) + obs_handle = open_obs_span(name, business_span_id=id, business_trace_id=self.trace_id) obs = obs_handle.correlation if obs_handle is not None else obs_correlation() if obs: serialized_data = {**(serialized_data or {}), **obs} diff --git a/tests/lib/core/tracing/test_obs_ids.py b/tests/lib/core/tracing/test_obs_ids.py index ddd079743..6f9418495 100644 --- a/tests/lib/core/tracing/test_obs_ids.py +++ b/tests/lib/core/tracing/test_obs_ids.py @@ -13,14 +13,14 @@ class TestGetObsMode: @pytest.mark.parametrize( "raw, expected", [ - (None, "dd_only"), # unset - ("", "dd_only"), # empty + (None, "dd_only"), # unset + ("", "dd_only"), # empty ("dd_only", "dd_only"), ("lgtm", "lgtm"), - ("LGTM", "lgtm"), # case-insensitive - (" lgtm ", "lgtm"), # trimmed - ("dual", "dd_only"), # removed mode -> safe degrade - ("garbage", "dd_only"), # unrecognized -> safe degrade + ("LGTM", "lgtm"), # case-insensitive + (" lgtm ", "lgtm"), # trimmed + ("dual", "dd_only"), # removed mode -> safe degrade + ("garbage", "dd_only"), # unrecognized -> safe degrade ], ) def test_mode_resolution(self, monkeypatch, raw, expected): @@ -36,9 +36,7 @@ def test_lgtm_mode_reads_otel_and_emits_underscored_keys(self, monkeypatch): monkeypatch.setenv("SGP_OBS_MODE", "lgtm") monkeypatch.setattr(obs_ids, "_lgtm_ids", lambda: ("otel_trace", "otel_span")) # In lgtm mode ddtrace must NOT be consulted. - monkeypatch.setattr( - obs_ids, "_ddtrace_ids", lambda: pytest.fail("ddtrace read in lgtm mode") - ) + monkeypatch.setattr(obs_ids, "_ddtrace_ids", lambda: pytest.fail("ddtrace read in lgtm mode")) assert obs_correlation() == { "obs_trace_id": "otel_trace", @@ -48,9 +46,7 @@ def test_lgtm_mode_reads_otel_and_emits_underscored_keys(self, monkeypatch): def test_dd_only_mode_reads_ddtrace(self, monkeypatch): monkeypatch.setenv("SGP_OBS_MODE", "dd_only") monkeypatch.setattr(obs_ids, "_ddtrace_ids", lambda: ("dd_trace", "dd_span")) - monkeypatch.setattr( - obs_ids, "_lgtm_ids", lambda: pytest.fail("otel read in dd_only mode") - ) + monkeypatch.setattr(obs_ids, "_lgtm_ids", lambda: pytest.fail("otel read in dd_only mode")) assert obs_correlation() == { "obs_trace_id": "dd_trace", @@ -61,9 +57,7 @@ def test_stale_dual_degrades_to_ddtrace(self, monkeypatch): """A leftover SGP_OBS_MODE=dual must behave as dd_only, not read OTel.""" monkeypatch.setenv("SGP_OBS_MODE", "dual") monkeypatch.setattr(obs_ids, "_ddtrace_ids", lambda: ("dd_trace", "dd_span")) - monkeypatch.setattr( - obs_ids, "_lgtm_ids", lambda: pytest.fail("otel read for stale dual mode") - ) + monkeypatch.setattr(obs_ids, "_lgtm_ids", lambda: pytest.fail("otel read for stale dual mode")) assert obs_correlation() == { "obs_trace_id": "dd_trace", diff --git a/tests/lib/core/tracing/test_obs_span.py b/tests/lib/core/tracing/test_obs_span.py index aa7d826bb..fd36fb2ba 100644 --- a/tests/lib/core/tracing/test_obs_span.py +++ b/tests/lib/core/tracing/test_obs_span.py @@ -6,8 +6,7 @@ import pytest -from agentex.lib.core.tracing import obs_span -from agentex.lib.core.tracing import trace as trace_module +from agentex.lib.core.tracing import trace as trace_module, obs_span from agentex.lib.core.tracing.trace import Trace @@ -135,13 +134,11 @@ def test_lgtm_opens_named_span_and_reads_its_ids(self, monkeypatch): monkeypatch.setenv("SGP_OBS_MODE", "lgtm") record = _install_fake_otel(monkeypatch, trace_id=0xABC, span_id=0xFF) - handle = obs_span.open_obs_span( - "rocket.tool.fetch", business_span_id="bspan-1", business_trace_id="btrace-1" - ) + handle = obs_span.open_obs_span("rocket.tool.fetch", business_span_id="bspan-1", business_trace_id="btrace-1") assert handle is not None - assert record["span"].name == "rocket.tool.fetch" # named for the step - assert len(record["attached"]) == 1 # made active + assert record["span"].name == "rocket.tool.fetch" # named for the step + assert len(record["attached"]) == 1 # made active assert handle.correlation == { "obs_trace_id": "00000000000000000000000000000abc", "obs_span_id": "000000000000000000ff"[-16:], @@ -161,9 +158,7 @@ def start_span(name): span._ctx = _FakeSpanContext(0, 0, is_valid=False) return span - sys.modules["opentelemetry"].trace.get_tracer = lambda _n: types.SimpleNamespace( - start_span=start_span - ) + sys.modules["opentelemetry"].trace.get_tracer = lambda _n: types.SimpleNamespace(start_span=start_span) handle = obs_span.open_obs_span("step") assert handle is not None assert handle.correlation == {} @@ -211,9 +206,7 @@ def test_dd_only_with_active_ctx_opens_named_span(self, monkeypatch): monkeypatch.setenv("SGP_OBS_MODE", "dd_only") record = _install_fake_ddtrace(monkeypatch, active=True, trace_id=0xABC, span_id=0xFF) - handle = obs_span.open_obs_span( - "rocket.tool.fetch", business_span_id="bspan-9", business_trace_id="btrace-9" - ) + handle = obs_span.open_obs_span("rocket.tool.fetch", business_span_id="bspan-9", business_trace_id="btrace-9") assert handle is not None assert record["span"].name == "rocket.tool.fetch" @@ -268,10 +261,10 @@ def test_lgtm_business_span_tagged_with_wrapper_ids(self, monkeypatch): trace = Trace(processors=[], client=MagicMock(), trace_id="task-run-1") span = trace.start_span(name="chat_completion") - assert record["span"].name == "chat_completion" # dedicated named span + assert record["span"].name == "chat_completion" # dedicated named span assert span.data["obs_trace_id"] == "00000000000000000000000000000111" assert span.data["obs_span_id"] == "0000000000000222" - assert span.trace_id == "task-run-1" # business id unchanged + assert span.trace_id == "task-run-1" # business id unchanged assert span.id in trace_module._OBS_HANDLES # bidirectional: the obs span carries the business ids (reverse tag), # and the business span carries the obs ids (forward edge). @@ -305,7 +298,7 @@ def test_wrapper_ends_across_separate_trace_instances(self, monkeypatch): ender = Trace(processors=[], client=MagicMock(), trace_id="task-run-x") ender.end_span(span) - assert record["span"].ended is True # wrapper WAS ended -> exportable + assert record["span"].ended is True # wrapper WAS ended -> exportable assert span.id not in trace_module._OBS_HANDLES # handle cleaned up def test_dd_only_business_span_tagged_via_ddtrace(self, monkeypatch): @@ -350,9 +343,9 @@ def test_dd_only_no_ctx_falls_back_to_ambient(self, monkeypatch): trace = Trace(processors=[], client=MagicMock(), trace_id="task-run-3") span = trace.start_span(name="get_state") - assert span.id not in trace_module._OBS_HANDLES # no wrapper opened - assert span.data is None # nothing tagged - trace.end_span(span) # must not raise + assert span.id not in trace_module._OBS_HANDLES # no wrapper opened + assert span.data is None # nothing tagged + trace.end_span(span) # must not raise # --------------------------------------------------------------------------- # @@ -366,8 +359,8 @@ def test_lgtm_touches_only_otel(self, monkeypatch): obs_span.open_obs_span("step") - assert otel["span"] is not None # OTel wrapper opened - assert dd["span"] is None # ddtrace never touched + assert otel["span"] is not None # OTel wrapper opened + assert dd["span"] is None # ddtrace never touched def test_dd_only_touches_only_ddtrace(self, monkeypatch): monkeypatch.setenv("SGP_OBS_MODE", "dd_only") @@ -376,8 +369,8 @@ def test_dd_only_touches_only_ddtrace(self, monkeypatch): obs_span.open_obs_span("step") - assert dd["span"] is not None # ddtrace wrapper opened - assert otel["span"] is None # OTel never touched + assert dd["span"] is not None # ddtrace wrapper opened + assert otel["span"] is None # OTel never touched # --------------------------------------------------------------------------- # @@ -402,7 +395,7 @@ def boom(_name): raise RuntimeError("tracer blew up") sys.modules["opentelemetry"].trace.get_tracer = boom - assert obs_span.open_obs_span("step") is None # inner guard + assert obs_span.open_obs_span("step") is None # inner guard def test_top_level_guard_swallows_get_mode_error(self, monkeypatch): # Even if mode resolution itself raises, open_obs_span must not. @@ -424,8 +417,8 @@ def test_unconfigured_lgtm_yields_usable_span_no_raise(self, monkeypatch): span = trace.start_span(name="safe") assert span.trace_id == "task-run-4" - assert span.id not in trace_module._OBS_HANDLES # no wrapper - trace.end_span(span) # must not raise + assert span.id not in trace_module._OBS_HANDLES # no wrapper + trace.end_span(span) # must not raise def _install_fake_otel_sequence(monkeypatch, *, trace_id: int, first_span_id: int): @@ -492,8 +485,8 @@ def test_turn2_each_step_gets_distinct_named_wrapper_under_trace_B(self, monkeyp for biz, wrapper, exp_span in zip(business, state["spans"], expected_obs_span): # forward edge: business span carries the wrapper's ids - assert biz.data["obs_trace_id"] == obs_trace_B # all under trace B - assert biz.data["obs_span_id"] == exp_span # distinct wBn + assert biz.data["obs_trace_id"] == obs_trace_B # all under trace B + assert biz.data["obs_span_id"] == exp_span # distinct wBn # reverse tag: wrapper carries the business ids assert wrapper.attributes == { "agentex.business_span_id": biz.id, From 2a0b694eef21f17818a0645f2fc4d6dec8606e20 Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Sun, 2 Aug 2026 22:41:22 -0700 Subject: [PATCH 06/12] fix(tracing): satisfy pyright in obs-correlation test mocks Annotate fake ModuleType stubs as Any (pyright rejects attribute assignment on ModuleType), widen the mock 'record' dicts to dict[str, Any], assert the Optional resolver returns before unpacking, and narrow span.data with isinstance before subscripting. Clears the pyright errors failing the lint job. No behavior change. Co-Authored-By: Claude Opus 4.8 --- tests/lib/core/tracing/test_obs_ids.py | 19 ++++++++++++------- tests/lib/core/tracing/test_obs_span.py | 15 +++++++++------ 2 files changed, 21 insertions(+), 13 deletions(-) diff --git a/tests/lib/core/tracing/test_obs_ids.py b/tests/lib/core/tracing/test_obs_ids.py index 6f9418495..5cdeb81b8 100644 --- a/tests/lib/core/tracing/test_obs_ids.py +++ b/tests/lib/core/tracing/test_obs_ids.py @@ -2,6 +2,7 @@ import sys import types +from typing import Any import pytest @@ -87,13 +88,15 @@ class TestIdFormatting: def test_ddtrace_ids_formats_w3c_hex(self, monkeypatch): ctx = types.SimpleNamespace(trace_id=0xABC, span_id=0xFF) tracer = types.SimpleNamespace(current_trace_context=lambda: ctx) - fake_ddtrace = types.ModuleType("ddtrace") - fake_trace = types.ModuleType("ddtrace.trace") + fake_ddtrace: Any = types.ModuleType("ddtrace") + fake_trace: Any = types.ModuleType("ddtrace.trace") fake_trace.tracer = tracer monkeypatch.setitem(sys.modules, "ddtrace", fake_ddtrace) monkeypatch.setitem(sys.modules, "ddtrace.trace", fake_trace) - trace_id, span_id = obs_ids._ddtrace_ids() + result = obs_ids._ddtrace_ids() + assert result is not None + trace_id, span_id = result assert trace_id == "00000000000000000000000000000abc" assert span_id == "000000000000000000ff"[-16:] # 16-hex assert len(trace_id) == 32 and len(span_id) == 16 @@ -102,18 +105,20 @@ def test_lgtm_ids_formats_w3c_hex(self, monkeypatch): span_ctx = types.SimpleNamespace(trace_id=0xABC, span_id=0xFF, is_valid=True) current_span = types.SimpleNamespace(get_span_context=lambda: span_ctx) fake_trace_mod = types.SimpleNamespace(get_current_span=lambda: current_span) - fake_otel = types.ModuleType("opentelemetry") + fake_otel: Any = types.ModuleType("opentelemetry") fake_otel.trace = fake_trace_mod monkeypatch.setitem(sys.modules, "opentelemetry", fake_otel) - trace_id, span_id = obs_ids._lgtm_ids() + result = obs_ids._lgtm_ids() + assert result is not None + trace_id, span_id = result assert trace_id == "00000000000000000000000000000abc" assert len(trace_id) == 32 and len(span_id) == 16 def test_ddtrace_ids_none_when_no_context(self, monkeypatch): tracer = types.SimpleNamespace(current_trace_context=lambda: None) - fake_ddtrace = types.ModuleType("ddtrace") - fake_trace = types.ModuleType("ddtrace.trace") + fake_ddtrace: Any = types.ModuleType("ddtrace") + fake_trace: Any = types.ModuleType("ddtrace.trace") fake_trace.tracer = tracer monkeypatch.setitem(sys.modules, "ddtrace", fake_ddtrace) monkeypatch.setitem(sys.modules, "ddtrace.trace", fake_trace) diff --git a/tests/lib/core/tracing/test_obs_span.py b/tests/lib/core/tracing/test_obs_span.py index fd36fb2ba..9a4ff69e4 100644 --- a/tests/lib/core/tracing/test_obs_span.py +++ b/tests/lib/core/tracing/test_obs_span.py @@ -2,6 +2,7 @@ import sys import types +from typing import Any from unittest.mock import MagicMock import pytest @@ -62,7 +63,7 @@ def end(self): def _install_fake_otel(monkeypatch, *, trace_id=0xABC, span_id=0xFF): - record: dict = {"span": None, "attached": [], "detached": []} + record: dict[str, Any] = {"span": None, "attached": [], "detached": []} def start_span(name): span = _FakeOtelSpan(name, trace_id, span_id) @@ -80,7 +81,7 @@ def start_span(name): attach=lambda ctx: record["attached"].append(ctx) or object(), detach=lambda token: record["detached"].append(token), ) - fake_otel = types.ModuleType("opentelemetry") + fake_otel: Any = types.ModuleType("opentelemetry") fake_otel.trace = fake_trace fake_otel.context = fake_context monkeypatch.setitem(sys.modules, "opentelemetry", fake_otel) @@ -104,7 +105,7 @@ def finish(self): def _install_fake_ddtrace(monkeypatch, *, active=True, trace_id=0xABC, span_id=0xFF): - record: dict = {"span": None, "started": []} + record: dict[str, Any] = {"span": None, "started": []} ctx_obj = object() if active else None record["ctx"] = ctx_obj @@ -118,8 +119,8 @@ def start_span(name, child_of=None, activate=False): current_trace_context=lambda: ctx_obj, start_span=start_span, ) - fake_ddtrace = types.ModuleType("ddtrace") - fake_trace = types.ModuleType("ddtrace.trace") + fake_ddtrace: Any = types.ModuleType("ddtrace") + fake_trace: Any = types.ModuleType("ddtrace.trace") fake_trace.tracer = tracer monkeypatch.setitem(sys.modules, "ddtrace", fake_ddtrace) monkeypatch.setitem(sys.modules, "ddtrace.trace", fake_trace) @@ -262,6 +263,7 @@ def test_lgtm_business_span_tagged_with_wrapper_ids(self, monkeypatch): span = trace.start_span(name="chat_completion") assert record["span"].name == "chat_completion" # dedicated named span + assert isinstance(span.data, dict) assert span.data["obs_trace_id"] == "00000000000000000000000000000111" assert span.data["obs_span_id"] == "0000000000000222" assert span.trace_id == "task-run-1" # business id unchanged @@ -309,6 +311,7 @@ def test_dd_only_business_span_tagged_via_ddtrace(self, monkeypatch): span = trace.start_span(name="get_state") assert record["span"].name == "get_state" + assert isinstance(span.data, dict) assert span.data["obs_trace_id"] == "00000000000000000000000000000111" assert span.data["obs_span_id"] == "0000000000000222" assert record["span"].tags == { @@ -444,7 +447,7 @@ def start_span(name): attach=lambda ctx: object(), detach=lambda token: None, ) - fake_otel = types.ModuleType("opentelemetry") + fake_otel: Any = types.ModuleType("opentelemetry") fake_otel.trace = fake_trace fake_otel.context = fake_context monkeypatch.setitem(sys.modules, "opentelemetry", fake_otel) From 1159d23186417eff1f8146a3b53f806f3b42d76d Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Tue, 4 Aug 2026 22:21:39 -0700 Subject: [PATCH 07/12] fix(tracing): skip obs wrapper inside Temporal activities start_span/end_span run as SEPARATE Temporal activities (START_SPAN/END_SPAN) that Temporal can route to different worker processes. The obs wrapper handle is stored in a process-local module dict, so on a multi-replica fleet the END lands on a different worker than the START: the handle is never popped (leak / OOM risk) and the wrapper span is never ended (dangling obs_span_id in Tempo). Inside a Temporal activity, skip opening our own wrapper and instead stamp the reverse tag onto the interceptor-propagated ambient span (tag_ambient_obs_span) and read forward ids via obs_correlation(). Trace-level correlation is preserved via the Temporal OTel TracingInterceptor (#485); the per-step named wrapper and TurnTrace RETRY/ASYNC roll-up are deferred (see TODO(obs-followup)). Co-Authored-By: Claude Opus 4.8 --- src/agentex/lib/core/tracing/obs_span.py | 39 +++++++++++++- src/agentex/lib/core/tracing/trace.py | 68 ++++++++++++++++++++++-- 2 files changed, 102 insertions(+), 5 deletions(-) diff --git a/src/agentex/lib/core/tracing/obs_span.py b/src/agentex/lib/core/tracing/obs_span.py index 37e01b430..541425d96 100644 --- a/src/agentex/lib/core/tracing/obs_span.py +++ b/src/agentex/lib/core/tracing/obs_span.py @@ -31,7 +31,7 @@ from agentex.lib.core.tracing.obs_ids import LGTM, get_obs_mode -__all__ = ("ObsSpanHandle", "open_obs_span", "close_obs_span") +__all__ = ("ObsSpanHandle", "open_obs_span", "close_obs_span", "tag_ambient_obs_span") # Instrumentation scope name so these wrapper spans are identifiable in Tempo/DD. _TRACER_NAME = "agentex.business" @@ -175,6 +175,43 @@ def open_obs_span( return None +def tag_ambient_obs_span( + business_span_id: Optional[str] = None, + business_trace_id: Optional[str] = None, +) -> 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 + ``agentex.business_span_id`` / ``agentex.business_trace_id`` so the obs -> business + pivot still works. Best-effort; never raises.""" + try: + if get_obs_mode() == LGTM: + from opentelemetry import trace + + span = trace.get_current_span() + if span is not None and span.get_span_context().is_valid: + if business_span_id: + span.set_attribute(_ATTR_BUSINESS_SPAN_ID, business_span_id) + if business_trace_id: + span.set_attribute(_ATTR_BUSINESS_TRACE_ID, business_trace_id) + else: + from ddtrace.trace import tracer + + span = tracer.current_span() + if span is not None: + if business_span_id: + span.set_tag(_ATTR_BUSINESS_SPAN_ID, business_span_id) + if business_trace_id: + span.set_tag(_ATTR_BUSINESS_TRACE_ID, business_trace_id) + except Exception: # pragma: no cover - best-effort; obs must never break a call + pass + + def close_obs_span( handle: Optional[ObsSpanHandle], error: Optional[Dict[str, str]] = None, diff --git a/src/agentex/lib/core/tracing/trace.py b/src/agentex/lib/core/tracing/trace.py index 0fddc1a98..d5419f02b 100644 --- a/src/agentex/lib/core/tracing/trace.py +++ b/src/agentex/lib/core/tracing/trace.py @@ -16,6 +16,7 @@ ObsSpanHandle, open_obs_span, close_obs_span, + tag_ambient_obs_span, ) from agentex.lib.core.tracing.span_error import get_span_error, set_span_error from agentex.lib.core.tracing.span_queue import ( @@ -42,6 +43,43 @@ _OBS_HANDLES: dict[str, ObsSpanHandle] = {} +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. + """ + try: + from temporalio import activity + + return activity.in_activity() + except Exception: + return False + + class Trace: """ Trace is a wrapper around the Agentex API for tracing. @@ -105,9 +143,20 @@ def start_span( # so you can pivot obs -> business in Tempo/DD. Falls back to the ambient # obs context (ddtrace) when not in lgtm mode. Business trace_id stays the # run-level task id. + # + # Inside a Temporal activity we skip the wrapper entirely and only tag the + # interceptor-propagated ambient span: opening a wrapper there would leak, + # since start_span / end_span run as separate activities on possibly + # different workers and the handle could never be closed. See + # _in_temporal_activity(). id = str(uuid.uuid4()) - obs_handle = open_obs_span(name, business_span_id=id, business_trace_id=self.trace_id) - obs = obs_handle.correlation if obs_handle is not None else obs_correlation() + if _in_temporal_activity(): + tag_ambient_obs_span(business_span_id=id, business_trace_id=self.trace_id) + obs_handle = None + obs = obs_correlation() + else: + obs_handle = open_obs_span(name, business_span_id=id, business_trace_id=self.trace_id) + obs = obs_handle.correlation if obs_handle is not None else obs_correlation() if obs: serialized_data = {**(serialized_data or {}), **obs} @@ -274,9 +323,20 @@ async def start_span( # so you can pivot obs -> business in Tempo/DD. Falls back to the ambient # obs context (ddtrace) when not in lgtm mode. Business trace_id stays the # run-level task id. + # + # Inside a Temporal activity we skip the wrapper entirely and only tag the + # interceptor-propagated ambient span: opening a wrapper there would leak, + # since start_span / end_span run as separate activities on possibly + # different workers and the handle could never be closed. See + # _in_temporal_activity(). id = str(uuid.uuid4()) - obs_handle = open_obs_span(name, business_span_id=id, business_trace_id=self.trace_id) - obs = obs_handle.correlation if obs_handle is not None else obs_correlation() + if _in_temporal_activity(): + tag_ambient_obs_span(business_span_id=id, business_trace_id=self.trace_id) + obs_handle = None + obs = obs_correlation() + else: + obs_handle = open_obs_span(name, business_span_id=id, business_trace_id=self.trace_id) + obs = obs_handle.correlation if obs_handle is not None else obs_correlation() if obs: serialized_data = {**(serialized_data or {}), **obs} From 044b3f4ed50c062b049ed7046360bfc4f62b916c Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Wed, 5 Aug 2026 00:13:50 -0700 Subject: [PATCH 08/12] feat(tracing): span the ACP->Temporal dispatch for end-to-end async traces The Temporal OTel interceptor propagates trace context by injecting the active span into the Temporal message headers on start_workflow/signal, so the worker roots the workflow+activity under it. But the ACP server dispatches from a bare async handler with no active span -> nothing injected -> the workflow's activity becomes a detached trace root, disconnected from the task/create / event/send that triggered it. Wrap submit_task and send_event in an OTel span (agentex.acp) so the interceptor has a context to inject. Child of the ingress request span when one is active (front-of-request propagation), else a per-turn root. Fail-open. Co-Authored-By: Claude Opus 4.8 --- .../services/temporal_task_service.py | 82 ++++++++++++------- 1 file changed, 54 insertions(+), 28 deletions(-) diff --git a/src/agentex/lib/core/temporal/services/temporal_task_service.py b/src/agentex/lib/core/temporal/services/temporal_task_service.py index 5f6c0c381..427779105 100644 --- a/src/agentex/lib/core/temporal/services/temporal_task_service.py +++ b/src/agentex/lib/core/temporal/services/temporal_task_service.py @@ -1,7 +1,8 @@ from __future__ import annotations -from typing import Any +from typing import Any, Iterator from datetime import timedelta +from contextlib import contextmanager from agentex.types.task import Task from agentex.types.agent import Agent @@ -13,6 +14,34 @@ from agentex.lib.core.clients.temporal.temporal_client import TemporalClient +@contextmanager +def _acp_dispatch_span(name: str) -> Iterator[None]: + """Wrap an ACP -> Temporal dispatch (start_workflow / signal) in an OTel span. + + The Temporal OpenTelemetry interceptor propagates trace context by injecting + the CURRENTLY ACTIVE span into the Temporal message headers on the caller + side (``start_workflow`` / ``signal_workflow``); the worker then extracts it + and roots the workflow / activity spans under it. But the ACP server dispatches + from a bare async handler with no active span, so nothing is injected and the + workflow's activities become DETACHED trace roots -- the business work shows up + in Tempo as a fresh trace with no link back to the ``task/create`` / + ``event/send`` that triggered it. + + Opening a span here gives the interceptor something to inject. It becomes a + child of the ingress request span when one is active (front-of-request + propagation), or a fresh per-turn root otherwise. Fail-open: never raises if + OpenTelemetry isn't importable. + """ + try: + from opentelemetry import trace as _otel_trace + except Exception: # pragma: no cover - obs must never break a dispatch + yield + return + tracer = _otel_trace.get_tracer("agentex.acp") + with tracer.start_as_current_span(name, kind=_otel_trace.SpanKind.PRODUCER): + yield + + class TemporalTaskService: """ Submits Agent agent_tasks to the async runtime for execution. @@ -26,7 +55,6 @@ def __init__( self._temporal_client = temporal_client self._env_vars = env_vars - async def submit_task(self, agent: Agent, task: Task, params: dict[str, Any] | None) -> str: """ Submit a task to the async runtime for execution. @@ -37,22 +65,19 @@ async def submit_task(self, agent: Agent, task: Task, params: dict[str, Any] | N # indefinitely, which long-lived chat/session agents rely on). A positive # value bounds the whole continue-as-new chain's wall-clock lifetime. timeout_seconds = self._env_vars.WORKFLOW_EXECUTION_TIMEOUT_SECONDS - execution_timeout = ( - timedelta(seconds=timeout_seconds) - if timeout_seconds and timeout_seconds > 0 - else None - ) - return await self._temporal_client.start_workflow( - workflow=self._env_vars.WORKFLOW_NAME, - arg=CreateTaskParams( - agent=agent, - task=task, - params=params, - ), - id=task.id, - task_queue=self._env_vars.WORKFLOW_TASK_QUEUE, - execution_timeout=execution_timeout, - ) + execution_timeout = timedelta(seconds=timeout_seconds) if timeout_seconds and timeout_seconds > 0 else None + with _acp_dispatch_span(f"acp.task_create:{task.id}"): + return await self._temporal_client.start_workflow( + workflow=self._env_vars.WORKFLOW_NAME, + arg=CreateTaskParams( + agent=agent, + task=task, + params=params, + ), + id=task.id, + task_queue=self._env_vars.WORKFLOW_TASK_QUEUE, + execution_timeout=execution_timeout, + ) async def get_state(self, task_id: str) -> WorkflowState: """ @@ -63,16 +88,17 @@ async def get_state(self, task_id: str) -> WorkflowState: ) async def send_event(self, agent: Agent, task: Task, event: Event, request: dict | None = None) -> None: - return await self._temporal_client.send_signal( - workflow_id=task.id, - signal=SignalName.RECEIVE_EVENT.value, - payload=SendEventParams( - agent=agent, - task=task, - event=event, - request=request, - ).model_dump(), - ) + with _acp_dispatch_span(f"acp.event_send:{task.id}"): + return await self._temporal_client.send_signal( + workflow_id=task.id, + signal=SignalName.RECEIVE_EVENT.value, + payload=SendEventParams( + agent=agent, + task=task, + event=event, + request=request, + ).model_dump(), + ) async def interrupt(self, agent: Agent, task: Task, request: dict | None = None) -> None: """Forward a task/interrupt to the running workflow as a dedicated signal. From 2740bac05e52a5481e28b16dabd7b80f7bba3199 Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Wed, 5 Aug 2026 22:26:13 -0700 Subject: [PATCH 09/12] refactor(tracing): low-cardinality ACP dispatch span names acp.task_create:{task.id} / acp.event_send:{task.id} put the task id in the span NAME, which is high-cardinality and breaks span-name aggregation in Tempo. Use static names (acp.task_create / acp.event_send) and carry the id as the agentex.task_id span attribute instead. Co-Authored-By: Claude Opus 4.8 --- .../core/temporal/services/temporal_task_service.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/agentex/lib/core/temporal/services/temporal_task_service.py b/src/agentex/lib/core/temporal/services/temporal_task_service.py index 427779105..2123b671d 100644 --- a/src/agentex/lib/core/temporal/services/temporal_task_service.py +++ b/src/agentex/lib/core/temporal/services/temporal_task_service.py @@ -15,7 +15,7 @@ @contextmanager -def _acp_dispatch_span(name: str) -> Iterator[None]: +def _acp_dispatch_span(name: str, task_id: str | None = None) -> Iterator[None]: """Wrap an ACP -> Temporal dispatch (start_workflow / signal) in an OTel span. The Temporal OpenTelemetry interceptor propagates trace context by injecting @@ -38,7 +38,10 @@ def _acp_dispatch_span(name: str) -> Iterator[None]: yield return tracer = _otel_trace.get_tracer("agentex.acp") - with tracer.start_as_current_span(name, kind=_otel_trace.SpanKind.PRODUCER): + # task_id goes on an attribute, NOT in the span name: a per-task span name is + # high-cardinality and breaks span-name aggregation in Tempo. + attributes = {"agentex.task_id": task_id} if task_id else None + with tracer.start_as_current_span(name, kind=_otel_trace.SpanKind.PRODUCER, attributes=attributes): yield @@ -66,7 +69,7 @@ async def submit_task(self, agent: Agent, task: Task, params: dict[str, Any] | N # value bounds the whole continue-as-new chain's wall-clock lifetime. timeout_seconds = self._env_vars.WORKFLOW_EXECUTION_TIMEOUT_SECONDS execution_timeout = timedelta(seconds=timeout_seconds) if timeout_seconds and timeout_seconds > 0 else None - with _acp_dispatch_span(f"acp.task_create:{task.id}"): + with _acp_dispatch_span("acp.task_create", task_id=task.id): return await self._temporal_client.start_workflow( workflow=self._env_vars.WORKFLOW_NAME, arg=CreateTaskParams( @@ -88,7 +91,7 @@ async def get_state(self, task_id: str) -> WorkflowState: ) async def send_event(self, agent: Agent, task: Task, event: Event, request: dict | None = None) -> None: - with _acp_dispatch_span(f"acp.event_send:{task.id}"): + with _acp_dispatch_span("acp.event_send", task_id=task.id): return await self._temporal_client.send_signal( workflow_id=task.id, signal=SignalName.RECEIVE_EVENT.value, From 9018a61b4f5e5cd8d050e88caf993d4642d4d7bc Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Wed, 5 Aug 2026 23:31:04 -0700 Subject: [PATCH 10/12] =?UTF-8?q?fix(tracing):=20address=20PR=20#484=20rev?= =?UTF-8?q?iew=20=E2=80=94=20error=20status,=20obs=20fallback,=20handle=20?= =?UTF-8?q?leak,=20dd=5Fonly=20Temporal=20backend?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four review-comment fixes for the business<->obs correlation edge, each with unit coverage: - ADK span() CM recorded no error on a failing step, so a failed agent step closed a green obs span. Now sets span error (guarded so obs can't shadow the app exception) and re-raises. (tests/test_adk_tracing_span_error.py) - In lgtm mode with no TracerProvider, open_obs_span returned a handle with empty correlation, suppressing the ambient obs_correlation() fallback -> business span got no obs_* ids. Now bails to None so the caller falls back. (tests/test_obs_span_fallback.py) - Obs-handle registry could leak (registration-order + start-without-end via public API). Processor hooks now swallow (obs must never crash the app path) and the registry is a bounded OrderedDict that evicts+closes the oldest. (tests/test_obs_handle_registry.py) - On the Temporal path the ambient span is the OTel interceptor span regardless of SGP_OBS_MODE, but tag/read branched on mode -> in the default dd_only they tagged/read an unrelated ddtrace span. Added prefer_otel (OTel-first, ddtrace fallback) via a shared _begin_obs helper used by both start_span paths. (tests/test_temporal_obs_backend.py) Also: ObsSpanHandle.close() instead of reaching into _close; Iterator from collections.abc. Co-Authored-By: Claude Opus 4.8 --- src/agentex/lib/adk/_modules/tracing.py | 19 ++ .../services/temporal_task_service.py | 3 +- src/agentex/lib/core/tracing/obs_ids.py | 12 +- src/agentex/lib/core/tracing/obs_span.py | 96 +++++++--- src/agentex/lib/core/tracing/trace.py | 166 +++++++++++++----- tests/test_adk_tracing_span_error.py | 108 ++++++++++++ tests/test_obs_handle_registry.py | 126 +++++++++++++ tests/test_obs_span_fallback.py | 116 ++++++++++++ tests/test_temporal_obs_backend.py | 134 ++++++++++++++ 9 files changed, 712 insertions(+), 68 deletions(-) create mode 100644 tests/test_adk_tracing_span_error.py create mode 100644 tests/test_obs_handle_registry.py create mode 100644 tests/test_obs_span_fallback.py create mode 100644 tests/test_temporal_obs_backend.py diff --git a/src/agentex/lib/adk/_modules/tracing.py b/src/agentex/lib/adk/_modules/tracing.py index 7d49bb91c..4a58be4e5 100644 --- a/src/agentex/lib/adk/_modules/tracing.py +++ b/src/agentex/lib/adk/_modules/tracing.py @@ -20,6 +20,7 @@ TracingActivityName, ) from agentex.lib.core.tracing.tracer import AsyncTracer +from agentex.lib.core.tracing.span_error import set_span_error from agentex.lib.core.harness.types import TurnUsage from agentex.types.span import Span from agentex.lib.utils.logging import make_logger @@ -236,6 +237,24 @@ async def span( ) try: yield span + except Exception as exc: + # Record the failure on the span so the obs span reflects the error + # instead of a false green. Agents use THIS context manager (not + # AsyncTrace.span, which is the only other place set_span_error is + # called), so without this a failed step closes green. end_span (in + # finally) reads it via get_span_error and propagates it to + # close_obs_span. Stored on span.data, so it round-trips through the + # END_SPAN activity on the Temporal path too. + # + # Guard set_span_error itself: it's obs work and must never replace + # the app's exception on the way out. We always re-raise the ORIGINAL + # exc regardless. + if span: + try: + set_span_error(span, exc) + except Exception: # pragma: no cover - obs must not break app path + pass + raise finally: if span: await self.end_span( diff --git a/src/agentex/lib/core/temporal/services/temporal_task_service.py b/src/agentex/lib/core/temporal/services/temporal_task_service.py index 2123b671d..d12f20eb8 100644 --- a/src/agentex/lib/core/temporal/services/temporal_task_service.py +++ b/src/agentex/lib/core/temporal/services/temporal_task_service.py @@ -1,8 +1,9 @@ from __future__ import annotations -from typing import Any, Iterator +from typing import Any from datetime import timedelta from contextlib import contextmanager +from collections.abc import Iterator from agentex.types.task import Task from agentex.types.agent import Agent diff --git a/src/agentex/lib/core/tracing/obs_ids.py b/src/agentex/lib/core/tracing/obs_ids.py index aa5b2678c..45fada783 100644 --- a/src/agentex/lib/core/tracing/obs_ids.py +++ b/src/agentex/lib/core/tracing/obs_ids.py @@ -69,7 +69,7 @@ def _ddtrace_ids() -> Optional[Tuple[str, str]]: return None -def obs_correlation() -> Dict[str, str]: +def obs_correlation(prefer_otel: bool = False) -> Dict[str, str]: """Return ``{"obs_trace_id": ..., "obs_span_id": ...}`` for the active observability context, or ``{}`` if none is active. @@ -79,10 +79,18 @@ def obs_correlation() -> 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 + ``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: - ids = _lgtm_ids() if get_obs_mode() == LGTM else _ddtrace_ids() + if prefer_otel: + ids = _lgtm_ids() or _ddtrace_ids() + else: + ids = _lgtm_ids() if get_obs_mode() == LGTM else _ddtrace_ids() except Exception: # obs must never fail an app call return {} diff --git a/src/agentex/lib/core/tracing/obs_span.py b/src/agentex/lib/core/tracing/obs_span.py index 541425d96..385507269 100644 --- a/src/agentex/lib/core/tracing/obs_span.py +++ b/src/agentex/lib/core/tracing/obs_span.py @@ -56,6 +56,11 @@ def __init__( self.correlation = correlation self._close = close + def close(self, error: Optional[Dict[str, str]] = None) -> None: + """Run the backend-specific closer (detach+end for OTel, finish for + ddtrace). ``error`` marks the obs span failed so it isn't a false green.""" + self._close(error) + def _hex_ids(trace_id: int, span_id: int) -> Dict[str, str]: """W3C-hex form: 32-hex trace, 16-hex span.""" @@ -82,7 +87,19 @@ def _open_otel_span( span.set_attribute(_ATTR_BUSINESS_TRACE_ID, business_trace_id) token = context.attach(trace.set_span_in_context(span)) sc = span.get_span_context() - correlation = _hex_ids(sc.trace_id, sc.span_id) if (sc and sc.is_valid) else {} + if not (sc and sc.is_valid): + # No real TracerProvider installed (lgtm mode but the agent has no + # OTel provider yet): the proxy tracer hands back a NonRecordingSpan + # with an invalid context. Returning a handle with empty correlation + # here would make the caller (trace.py) take obs_handle.correlation + # == {} and NEVER consult the obs_correlation() ambient fallback -- + # so the business span would get no obs_* ids at all, strictly worse + # than falling back. Detach the useless context, end the no-op span, + # and return None so the caller uses the ambient ids instead. + context.detach(token) + span.end() + return None + correlation = _hex_ids(sc.trace_id, sc.span_id) def _close(error: Optional[Dict[str, str]] = None) -> None: try: @@ -124,11 +141,18 @@ def _open_ddtrace_span( # Datadog traces. Parenting to the active request/turn context rolls them # into one trace while obs_span_id stays distinct per step. span = tracer.start_span(name, child_of=ctx, activate=True) + if not span.trace_id: + # Symmetry with the OTel path: a handle carrying empty correlation + # would suppress the ambient obs_correlation() fallback in trace.py. + # (child_of=ctx normally guarantees a real trace_id, so this is + # belt-and-braces.) Finish the span and fall back to ambient ids. + span.finish() + return None if business_span_id: span.set_tag(_ATTR_BUSINESS_SPAN_ID, business_span_id) if business_trace_id: span.set_tag(_ATTR_BUSINESS_TRACE_ID, business_trace_id) - correlation = _hex_ids(span.trace_id, span.span_id) if span.trace_id else {} + correlation = _hex_ids(span.trace_id, span.span_id) def _close(error: Optional[Dict[str, str]] = None) -> None: try: @@ -175,9 +199,44 @@ def open_obs_span( return None +def _tag_otel_ambient(business_span_id: Optional[str], business_trace_id: Optional[str]) -> bool: + """Stamp the reverse tag onto the active OTel span. Returns True iff a valid + OTel span was found and tagged.""" + try: + from opentelemetry import trace + except ImportError: + return False + span = trace.get_current_span() + if span is not None and span.get_span_context().is_valid: + if business_span_id: + span.set_attribute(_ATTR_BUSINESS_SPAN_ID, business_span_id) + if business_trace_id: + span.set_attribute(_ATTR_BUSINESS_TRACE_ID, business_trace_id) + return True + return False + + +def _tag_ddtrace_ambient(business_span_id: Optional[str], business_trace_id: Optional[str]) -> bool: + """Stamp the reverse tag onto the active ddtrace span. Returns True iff a + ddtrace span was found and tagged.""" + try: + from ddtrace.trace import tracer + except ImportError: + return False + span = tracer.current_span() + if span is not None: + if business_span_id: + span.set_tag(_ATTR_BUSINESS_SPAN_ID, business_span_id) + if business_trace_id: + span.set_tag(_ATTR_BUSINESS_TRACE_ID, business_trace_id) + return True + return False + + def tag_ambient_obs_span( business_span_id: Optional[str] = None, business_trace_id: Optional[str] = None, + prefer_otel: bool = False, ) -> None: """Stamp the reverse tag onto the CURRENTLY ACTIVE obs span -- without opening a new one. @@ -188,26 +247,23 @@ def tag_ambient_obs_span( 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.""" + pivot still works. Best-effort; never raises. + + ``prefer_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 + 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 _tag_otel_ambient(business_span_id, business_trace_id): + return + _tag_ddtrace_ambient(business_span_id, business_trace_id) + return if get_obs_mode() == LGTM: - from opentelemetry import trace - - span = trace.get_current_span() - if span is not None and span.get_span_context().is_valid: - if business_span_id: - span.set_attribute(_ATTR_BUSINESS_SPAN_ID, business_span_id) - if business_trace_id: - span.set_attribute(_ATTR_BUSINESS_TRACE_ID, business_trace_id) + _tag_otel_ambient(business_span_id, business_trace_id) else: - from ddtrace.trace import tracer - - span = tracer.current_span() - if span is not None: - if business_span_id: - span.set_tag(_ATTR_BUSINESS_SPAN_ID, business_span_id) - if business_trace_id: - span.set_tag(_ATTR_BUSINESS_TRACE_ID, business_trace_id) + _tag_ddtrace_ambient(business_span_id, business_trace_id) except Exception: # pragma: no cover - best-effort; obs must never break a call pass @@ -222,6 +278,6 @@ def close_obs_span( if handle is None: return try: - handle._close(error) + handle.close(error) except Exception: # pragma: no cover - best-effort pass diff --git a/src/agentex/lib/core/tracing/trace.py b/src/agentex/lib/core/tracing/trace.py index d5419f02b..d3decdb9b 100644 --- a/src/agentex/lib/core/tracing/trace.py +++ b/src/agentex/lib/core/tracing/trace.py @@ -4,6 +4,7 @@ from typing import Any, AsyncGenerator from datetime import UTC, datetime from contextlib import contextmanager, asynccontextmanager +from collections import OrderedDict from pydantic import BaseModel @@ -40,7 +41,69 @@ # span is never .end()ed -> never exported (Simple/Batch processors only emit on # end). A module-level dict keyed by the unique span id survives across instances; # uuid4 span ids cannot collide across concurrent traces. -_OBS_HANDLES: dict[str, ObsSpanHandle] = {} +# +# Bounded (OrderedDict + cap): a correct start_span/end_span pair pops its own +# entry, so the registry normally hovers near the live-span count. The cap only +# bites when a caller starts a span and never ends it -- adk.tracing.start_span / +# end_span are public, unpaired API, so a caller-side bug (crash / early return +# between start and end) would otherwise grow this unbounded in a long-lived ACP +# process. Past the cap we evict+close the OLDEST handle so the leak degrades +# gracefully instead of OOMing (and the evicted span still .end()s -> exports). +_OBS_HANDLES_MAX = 2048 +_OBS_HANDLES: OrderedDict[str, ObsSpanHandle] = OrderedDict() + + +def _register_obs_handle(span_id: str, handle: ObsSpanHandle) -> None: + """Register an open obs wrapper handle, bounding the registry at + ``_OBS_HANDLES_MAX``. When over the cap, evict and close the oldest handle + first. close_obs_span is best-effort (detach may warn since it runs on a + different stack than the attach) and always .end()s the span, so an evicted + span still exports rather than dangling.""" + _OBS_HANDLES[span_id] = handle + _OBS_HANDLES.move_to_end(span_id) + while len(_OBS_HANDLES) > _OBS_HANDLES_MAX: + _evicted_id, evicted = _OBS_HANDLES.popitem(last=False) + logger.warning( + "obs handle registry over cap (%d); evicting+closing oldest span %r. " + "This means a caller started a span without ending it.", + _OBS_HANDLES_MAX, + _evicted_id, + ) + close_obs_span(evicted) + + +def _run_on_span_start(processor: SyncTracingProcessor, span: Span) -> None: + """Invoke ``on_span_start`` such that a processor bug can NEVER crash the app. + + Observability must degrade, not propagate: if this raised, the caller's + start_span would never return, the caller would never end_span, and the obs + handle would leak (dict entry + attached OTel context + unended span). By + swallowing here, start_span returns normally and the standard end_span path + pops and closes the handle -- no leak, no app-path failure.""" + try: + processor.on_span_start(span) + except Exception: + logger.warning( + "on_span_start raised for processor %r; skipping (observability must not fail the app path)", + type(processor).__name__, + exc_info=True, + ) + + +def _run_on_span_end(processor: SyncTracingProcessor, span: Span) -> None: + """Invoke ``on_span_end`` such that a processor bug can NEVER crash the app. + + Symmetric with :func:`_run_on_span_start`. The obs wrapper is already closed + before this runs (see end_span), so this only guards the app path against a + buggy processor -- there is no handle left to leak here.""" + try: + processor.on_span_end(span) + except Exception: + logger.warning( + "on_span_end raised for processor %r; skipping (observability must not fail the app path)", + type(processor).__name__, + exc_info=True, + ) def _in_temporal_activity() -> bool: @@ -80,6 +143,35 @@ def _in_temporal_activity() -> bool: return False +def _begin_obs( + name: str, + 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)``. + + 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``. + """ + 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() + return handle, correlation + + class Trace: """ Trace is a wrapper around the Agentex API for tracing. @@ -137,26 +229,10 @@ def start_span( serialized_input = recursive_model_dump(input) if input else None serialized_data = recursive_model_dump(data) if data else None - # Open a dedicated obs wrapper span named for this step and make it - # active, so obs_span_id is stable/meaningful (not an arbitrary innermost - # httpx span). It also carries the reverse tag (business span/trace id) - # so you can pivot obs -> business in Tempo/DD. Falls back to the ambient - # obs context (ddtrace) when not in lgtm mode. Business trace_id stays the - # run-level task id. - # - # Inside a Temporal activity we skip the wrapper entirely and only tag the - # interceptor-propagated ambient span: opening a wrapper there would leak, - # since start_span / end_span run as separate activities on possibly - # different workers and the handle could never be closed. See - # _in_temporal_activity(). + # Open the obs wrapper (or tag the ambient Temporal-activity span); see + # _begin_obs. Business trace_id stays the run-level task id. id = str(uuid.uuid4()) - if _in_temporal_activity(): - tag_ambient_obs_span(business_span_id=id, business_trace_id=self.trace_id) - obs_handle = None - obs = obs_correlation() - else: - obs_handle = open_obs_span(name, business_span_id=id, business_trace_id=self.trace_id) - obs = obs_handle.correlation if obs_handle is not None else obs_correlation() + obs_handle, obs = _begin_obs(name, id, self.trace_id) if obs: serialized_data = {**(serialized_data or {}), **obs} @@ -171,10 +247,10 @@ def start_span( task_id=task_id, ) if obs_handle is not None: - _OBS_HANDLES[span.id] = obs_handle + _register_obs_handle(span.id, obs_handle) for processor in self.processors: - processor.on_span_start(span) + _run_on_span_start(processor, span) return span @@ -203,7 +279,7 @@ def end_span( span.data = recursive_model_dump(span.data) if span.data else None for processor in self.processors: - processor.on_span_end(span) + _run_on_span_end(processor, span) return span @@ -317,26 +393,10 @@ async def start_span( serialized_input = recursive_model_dump(input) if input else None serialized_data = recursive_model_dump(data) if data else None - # Open a dedicated obs wrapper span named for this step and make it - # active, so obs_span_id is stable/meaningful (not an arbitrary innermost - # httpx span). It also carries the reverse tag (business span/trace id) - # so you can pivot obs -> business in Tempo/DD. Falls back to the ambient - # obs context (ddtrace) when not in lgtm mode. Business trace_id stays the - # run-level task id. - # - # Inside a Temporal activity we skip the wrapper entirely and only tag the - # interceptor-propagated ambient span: opening a wrapper there would leak, - # since start_span / end_span run as separate activities on possibly - # different workers and the handle could never be closed. See - # _in_temporal_activity(). + # Open the obs wrapper (or tag the ambient Temporal-activity span); see + # _begin_obs. Business trace_id stays the run-level task id. id = str(uuid.uuid4()) - if _in_temporal_activity(): - tag_ambient_obs_span(business_span_id=id, business_trace_id=self.trace_id) - obs_handle = None - obs = obs_correlation() - else: - obs_handle = open_obs_span(name, business_span_id=id, business_trace_id=self.trace_id) - obs = obs_handle.correlation if obs_handle is not None else obs_correlation() + obs_handle, obs = _begin_obs(name, id, self.trace_id) if obs: serialized_data = {**(serialized_data or {}), **obs} @@ -351,10 +411,20 @@ async def start_span( task_id=task_id, ) if obs_handle is not None: - _OBS_HANDLES[span.id] = obs_handle + _register_obs_handle(span.id, obs_handle) + # Enqueueing the START event must not crash the app path either (same + # principle as _run_on_span_start): swallow so start_span still returns + # and end_span cleans up the handle. The processors' on_span_start runs + # later on the queue worker, off the request path. if self.processors: - self._span_queue.enqueue(SpanEventType.START, span.model_copy(deep=True), self.processors) + try: + self._span_queue.enqueue(SpanEventType.START, span.model_copy(deep=True), self.processors) + except Exception: + logger.warning( + "failed to enqueue START span event; skipping (observability must not fail the app path)", + exc_info=True, + ) return span @@ -383,7 +453,13 @@ async def end_span( span.data = recursive_model_dump(span.data) if span.data else None if self.processors: - self._span_queue.enqueue(SpanEventType.END, span.model_copy(deep=True), self.processors) + try: + self._span_queue.enqueue(SpanEventType.END, span.model_copy(deep=True), self.processors) + except Exception: + logger.warning( + "failed to enqueue END span event; skipping (observability must not fail the app path)", + exc_info=True, + ) return span diff --git a/tests/test_adk_tracing_span_error.py b/tests/test_adk_tracing_span_error.py new file mode 100644 index 000000000..643115a82 --- /dev/null +++ b/tests/test_adk_tracing_span_error.py @@ -0,0 +1,108 @@ +"""Tests for the ADK ``TracingModule.span`` / ``turn_span`` error-status behavior. + +Regression coverage for the "false green" bug: agents open spans through the ADK +context manager (``adk.tracing.span`` / ``turn_span``), which is the *only* span +path they use. Before the fix, a failing step still closed its span green because +the CM never recorded the exception. These tests assert that: + + - a body exception is recorded on the span (``set_span_error`` -> ``data["__error__"]``), + - the ORIGINAL app exception always propagates unchanged, + - ``end_span`` sees the span *with* the error already set (except-before-finally), + - obs bookkeeping never breaks the app path (if ``set_span_error`` itself raises, + the app exception still propagates), + - the success path records no error, + - a falsy ``trace_id`` is a pure no-op (no start/end, yields ``None``), + - ``turn_span`` inherits all of the above since it delegates to ``span``. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock + +import pytest + +from agentex.types.span import Span +from agentex.lib.adk._modules.tracing import TracingModule +from agentex.lib.core.tracing.span_error import get_span_error + + +def _make_module() -> tuple[TracingModule, Span, AsyncMock]: + """A TracingModule with start_span/end_span stubbed to avoid any network. + + start_span returns a fresh Span; end_span is an AsyncMock so tests can + inspect the span (and its recorded error) as end_span actually saw it. + """ + module = TracingModule() + span = Span(id="span-1", name="step", start_time=1.0, trace_id="trace-1") + module.start_span = AsyncMock(return_value=span) # type: ignore[method-assign] + module.end_span = AsyncMock(return_value=span) # type: ignore[method-assign] + return module, span, module.end_span # type: ignore[return-value] + + +async def test_span_records_error_and_reraises() -> None: + module, span, end_span = _make_module() + + with pytest.raises(ValueError, match="boom"): + async with module.span(trace_id="trace-1", name="step") as yielded: + assert yielded is span + raise ValueError("boom") + + error = get_span_error(span) + assert error == {"type": "ValueError", "message": "boom"} + + # end_span still ran (finally) and saw the span with the error already set, + # so the failure is what gets persisted -- not a false green. + end_span.assert_awaited_once() + persisted_span = end_span.await_args.kwargs["span"] + assert get_span_error(persisted_span) == {"type": "ValueError", "message": "boom"} + + +async def test_span_success_records_no_error() -> None: + module, span, end_span = _make_module() + + async with module.span(trace_id="trace-1", name="step") as yielded: + assert yielded is span + + assert get_span_error(span) is None + end_span.assert_awaited_once() + + +async def test_span_obs_failure_does_not_shadow_app_exception(monkeypatch: pytest.MonkeyPatch) -> None: + """If set_span_error itself blows up, the app's exception must still surface.""" + module, span, end_span = _make_module() + + def _boom(*_args: object, **_kwargs: object) -> None: + raise RuntimeError("set_span_error is broken") + + monkeypatch.setattr("agentex.lib.adk._modules.tracing.set_span_error", _boom) + + # The ORIGINAL ValueError propagates, not the RuntimeError from obs code. + with pytest.raises(ValueError, match="boom"): + async with module.span(trace_id="trace-1", name="step"): + raise ValueError("boom") + + # The span still gets closed despite the obs hiccup. + end_span.assert_awaited_once() + + +async def test_span_noop_when_trace_id_falsy() -> None: + module, _span, end_span = _make_module() + + async with module.span(trace_id="", name="step") as yielded: + assert yielded is None + + module.start_span.assert_not_awaited() # type: ignore[attr-defined] + end_span.assert_not_awaited() + + +async def test_turn_span_records_error_and_reraises() -> None: + """turn_span delegates to span(), so it must record errors too.""" + module, span, end_span = _make_module() + + with pytest.raises(ValueError, match="boom"): + async with module.turn_span(trace_id="trace-1", name="turn") as turn: + assert turn.span is span + raise ValueError("boom") + + assert get_span_error(span) == {"type": "ValueError", "message": "boom"} + end_span.assert_awaited_once() diff --git a/tests/test_obs_handle_registry.py b/tests/test_obs_handle_registry.py new file mode 100644 index 000000000..02d3adf0a --- /dev/null +++ b/tests/test_obs_handle_registry.py @@ -0,0 +1,126 @@ +"""Tests for the obs-handle registry: leak safety + app-path safety. + +Two guarantees are pinned here: + + 1. A tracing processor whose ``on_span_start`` / ``on_span_end`` raises must + NOT crash the app path (``start_span`` / ``end_span`` still return). Because + start_span returns normally, the standard end_span path still pops+closes + the obs handle -- so the registration-order leak Greptile flagged cannot + happen. + 2. ``_OBS_HANDLES`` is bounded: a caller that starts spans without ending them + (public, unpaired ``start_span`` / ``end_span`` API) degrades gracefully -- + the oldest handle is evicted AND closed rather than growing unbounded. +""" + +from __future__ import annotations + +from typing import Any, cast + +import pytest +from opentelemetry import trace as otel_trace +from opentelemetry.trace import ( + TraceFlags, + SpanContext, + NonRecordingSpan, +) + +import agentex.lib.core.tracing.trace as trace_mod +from agentex.types.span import Span +from agentex.lib.core.tracing.trace import _OBS_HANDLES, _OBS_HANDLES_MAX, Trace +from agentex.lib.core.tracing.obs_span import ObsSpanHandle + + +@pytest.fixture(autouse=True) +def _clear_registry() -> Any: + """The registry is module-level global; keep tests isolated.""" + _OBS_HANDLES.clear() + yield + _OBS_HANDLES.clear() + + +def _valid_wrapper_span() -> NonRecordingSpan: + ctx = SpanContext( + trace_id=0x0123456789ABCDEF0123456789ABCDEF, + span_id=0x0123456789ABCDEF, + is_remote=False, + trace_flags=TraceFlags(TraceFlags.SAMPLED), + ) + return NonRecordingSpan(ctx) + + +class _RaisingProcessor: + """A processor whose lifecycle hooks blow up -- an obs bug must not crash the app.""" + + def __init__(self) -> None: + self.started = 0 + self.ended = 0 + + def on_span_start(self, span: Span) -> None: + self.started += 1 + raise RuntimeError("processor on_span_start is broken") + + def on_span_end(self, span: Span) -> None: + self.ended += 1 + raise RuntimeError("processor on_span_end is broken") + + +def _trace_with(processors: list[Any]) -> Trace: + return Trace(processors=processors, client=cast(Any, object()), trace_id="trace-1") + + +def test_start_span_survives_raising_processor_and_no_leak(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("SGP_OBS_MODE", "lgtm") + # Wrapper opens with a valid context -> a real handle is registered. + monkeypatch.setattr( + otel_trace, + "get_tracer", + lambda *a, **k: type("T", (), {"start_span": staticmethod(lambda *a, **k: _valid_wrapper_span())})(), + ) + + proc = _RaisingProcessor() + trace_obj = _trace_with([proc]) + + # A processor exploding in on_span_start must NOT propagate. + span = trace_obj.start_span(name="step") + assert proc.started == 1 + # The handle was registered despite the processor blowing up afterwards. + assert span.id in _OBS_HANDLES + + # end_span also survives a raising on_span_end AND pops/closes the handle, + # so nothing leaks. + trace_obj.end_span(span) + assert proc.ended == 1 + assert span.id not in _OBS_HANDLES + + +def test_registry_is_bounded_and_evicts_and_closes_oldest() -> None: + closed: list[str] = [] + + def _make_handle(marker: str) -> ObsSpanHandle: + return ObsSpanHandle(correlation={}, close=lambda _err=None, _m=marker: closed.append(_m)) + + # Fill exactly to the cap: nothing evicted yet. + for i in range(_OBS_HANDLES_MAX): + trace_mod._register_obs_handle(f"span-{i}", _make_handle(f"span-{i}")) + assert len(_OBS_HANDLES) == _OBS_HANDLES_MAX + assert closed == [] + + # One over the cap: the OLDEST (span-0) is evicted AND closed. + trace_mod._register_obs_handle("span-overflow", _make_handle("span-overflow")) + assert len(_OBS_HANDLES) == _OBS_HANDLES_MAX + assert "span-0" not in _OBS_HANDLES + assert "span-overflow" in _OBS_HANDLES + assert closed == ["span-0"] # evicted handle was closed, not just dropped + + +def test_reinserting_same_span_id_refreshes_recency() -> None: + def _noop_handle() -> ObsSpanHandle: + return ObsSpanHandle(correlation={}, close=lambda _err=None: None) + + trace_mod._register_obs_handle("a", _noop_handle()) + trace_mod._register_obs_handle("b", _noop_handle()) + # Touch "a" again -> it becomes the most-recent, so "b" is now the oldest. + trace_mod._register_obs_handle("a", _noop_handle()) + + oldest_key = next(iter(_OBS_HANDLES)) + assert oldest_key == "b" diff --git a/tests/test_obs_span_fallback.py b/tests/test_obs_span_fallback.py new file mode 100644 index 000000000..c92a42e34 --- /dev/null +++ b/tests/test_obs_span_fallback.py @@ -0,0 +1,116 @@ +"""Tests for the obs-wrapper -> ambient-correlation fallback. + +Regression coverage for: in ``lgtm`` mode with no OTel TracerProvider installed +(the documented current state of agents), ``open_obs_span`` used to return a +handle carrying an *empty* correlation. At the call site (``trace.py``) that +handle is not None, so the ambient ``obs_correlation()`` fallback was never +consulted and the business span ended up with **no** ``obs_*`` ids at all -- +strictly worse than falling back. + +The fix: ``open_obs_span`` bails out to ``None`` when the wrapper span's context +is invalid (proxy ``NonRecordingSpan``), so the caller falls back to the ambient +obs ids. These tests pin: + + - invalid wrapper context -> ``open_obs_span`` returns ``None`` and restores + the active context (no leaked attach), + - valid wrapper context -> a handle with real 32/16-hex correlation, + - end-to-end: with an invalid wrapper but a valid *ambient* span active, + ``Trace.start_span`` stamps the ambient ``obs_trace_id`` / ``obs_span_id`` + onto the business span (the fallback fires). +""" + +from __future__ import annotations + +from typing import Any, cast + +import pytest +from opentelemetry import trace as otel_trace, context as otel_context +from opentelemetry.trace import ( + INVALID_SPAN_CONTEXT, + TraceFlags, + SpanContext, + NonRecordingSpan, + set_span_in_context, +) + +from agentex.lib.core.tracing.trace import Trace +from agentex.lib.core.tracing.obs_span import open_obs_span, close_obs_span + +# Deterministic, valid ids for the "provider present" / ambient-span cases. +_TRACE_ID = 0x0123456789ABCDEF0123456789ABCDEF +_SPAN_ID = 0x0123456789ABCDEF +_TRACE_HEX = format(_TRACE_ID, "032x") +_SPAN_HEX = format(_SPAN_ID, "016x") + + +def _valid_span() -> NonRecordingSpan: + ctx = SpanContext( + trace_id=_TRACE_ID, + span_id=_SPAN_ID, + is_remote=False, + trace_flags=TraceFlags(TraceFlags.SAMPLED), + ) + return NonRecordingSpan(ctx) + + +class _FakeTracer: + """A tracer whose start_span returns a fixed span (bypasses any real provider).""" + + def __init__(self, span: NonRecordingSpan): + self._span = span + + def start_span(self, name: str, *args: object, **kwargs: object) -> NonRecordingSpan: + return self._span + + +def _patch_wrapper_tracer(monkeypatch: pytest.MonkeyPatch, span: NonRecordingSpan) -> None: + """Force the obs wrapper's ``trace.get_tracer(...).start_span`` to yield ``span``. + + Only affects the wrapper opened inside open_obs_span; obs_correlation reads + the *current* span via ``trace.get_current_span()`` and is untouched. + """ + monkeypatch.setattr(otel_trace, "get_tracer", lambda *a, **k: _FakeTracer(span)) + + +def test_open_obs_span_returns_none_on_invalid_context(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("SGP_OBS_MODE", "lgtm") + _patch_wrapper_tracer(monkeypatch, NonRecordingSpan(INVALID_SPAN_CONTEXT)) + + before = otel_trace.get_current_span() + handle = open_obs_span("step", business_span_id="bs", business_trace_id="bt") + + # No handle -> caller falls back to obs_correlation() instead of an empty {}. + assert handle is None + # The context attach inside open_obs_span was detached: no leak. + assert otel_trace.get_current_span() is before + + +def test_open_obs_span_returns_handle_on_valid_context(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("SGP_OBS_MODE", "lgtm") + _patch_wrapper_tracer(monkeypatch, _valid_span()) + + handle = open_obs_span("step", business_span_id="bs", business_trace_id="bt") + + assert handle is not None + assert handle.correlation == {"obs_trace_id": _TRACE_HEX, "obs_span_id": _SPAN_HEX} + close_obs_span(handle) + + +def test_start_span_falls_back_to_ambient_when_wrapper_invalid(monkeypatch: pytest.MonkeyPatch) -> None: + """End-to-end: invalid wrapper -> ambient obs ids land on the business span.""" + monkeypatch.setenv("SGP_OBS_MODE", "lgtm") + # Wrapper span has an invalid context (no real provider) -> open_obs_span None. + _patch_wrapper_tracer(monkeypatch, NonRecordingSpan(INVALID_SPAN_CONTEXT)) + + # But a VALID ambient span is active (e.g. the ACP ingress / interceptor span). + token = otel_context.attach(set_span_in_context(_valid_span())) + try: + trace_obj = Trace(processors=[], client=cast(Any, object()), trace_id="trace-1") + span = trace_obj.start_span(name="step") + finally: + otel_context.detach(token) + + # obs_correlation() was consulted and stamped the ambient ids onto data. + assert isinstance(span.data, dict) + assert span.data["obs_trace_id"] == _TRACE_HEX + assert span.data["obs_span_id"] == _SPAN_HEX diff --git a/tests/test_temporal_obs_backend.py b/tests/test_temporal_obs_backend.py new file mode 100644 index 000000000..34daf1d11 --- /dev/null +++ b/tests/test_temporal_obs_backend.py @@ -0,0 +1,134 @@ +"""Tests for the Temporal-path obs backend selection. + +Inside a Temporal activity the ambient span is temporalio's OpenTelemetry +``TracingInterceptor`` span -- always OTel, regardless of ``SGP_OBS_MODE``. The +reverse tag (``tag_ambient_obs_span``) and the forward correlation read +(``obs_correlation``) must therefore target OTel there, even in the default +``dd_only`` mode. Before the fix they branched on ``SGP_OBS_MODE`` and, in +``dd_only``, tagged/read an unrelated ddtrace span -- so the business<->obs +correlation on the async/Temporal path pointed at the wrong trace (or nowhere). +""" + +from __future__ import annotations + +from typing import Any, cast + +import pytest +from opentelemetry import trace as otel_trace +from opentelemetry.trace import TraceFlags, SpanContext + +import agentex.lib.core.tracing.trace as trace_mod +import agentex.lib.core.tracing.obs_ids as obs_ids_mod +from agentex.lib.core.tracing.trace import _OBS_HANDLES, Trace +from agentex.lib.core.tracing.obs_ids import obs_correlation +from agentex.lib.core.tracing.obs_span import tag_ambient_obs_span + +_TRACE_ID = 0x0123456789ABCDEF0123456789ABCDEF +_SPAN_ID = 0x0123456789ABCDEF +_TRACE_HEX = format(_TRACE_ID, "032x") +_SPAN_HEX = format(_SPAN_ID, "016x") + + +def _valid_ctx() -> SpanContext: + return SpanContext( + trace_id=_TRACE_ID, + span_id=_SPAN_ID, + is_remote=True, # like a Temporal-propagated remote parent + trace_flags=TraceFlags(TraceFlags.SAMPLED), + ) + + +class _RecordingOtelSpan: + """A stand-in for the interceptor's activity span that records set_attribute.""" + + def __init__(self, ctx: SpanContext) -> None: + self._ctx = ctx + self.attributes: dict[str, Any] = {} + + def get_span_context(self) -> SpanContext: + return self._ctx + + def set_attribute(self, key: str, value: Any) -> None: + self.attributes[key] = value + + +@pytest.fixture(autouse=True) +def _clear_registry() -> Any: + _OBS_HANDLES.clear() + yield + _OBS_HANDLES.clear() + + +def _activate_otel_span(monkeypatch: pytest.MonkeyPatch) -> _RecordingOtelSpan: + span = _RecordingOtelSpan(_valid_ctx()) + monkeypatch.setattr(otel_trace, "get_current_span", lambda *a, **k: span) + return span + + +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. + monkeypatch.setenv("SGP_OBS_MODE", "dd_only") + monkeypatch.setattr(trace_mod, "_in_temporal_activity", lambda: True) + activity_span = _activate_otel_span(monkeypatch) + + trace_obj = Trace(processors=[], client=cast(Any, object()), trace_id="trace-1") + span = trace_obj.start_span(name="process_turn") + + # Reverse tag landed on the OTel activity span (not a ddtrace span / nowhere). + assert activity_span.attributes["agentex.business_span_id"] == span.id + assert activity_span.attributes["agentex.business_trace_id"] == "trace-1" + + # Forward correlation recorded the OTel activity trace ids. + assert isinstance(span.data, dict) + assert span.data["obs_trace_id"] == _TRACE_HEX + assert span.data["obs_span_id"] == _SPAN_HEX + + # Temporal path opens no wrapper -> no handle registered (nothing to leak). + assert span.id not in _OBS_HANDLES + + +def test_obs_correlation_prefer_otel_prefers_otel_over_ddtrace(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("SGP_OBS_MODE", "dd_only") + _activate_otel_span(monkeypatch) + # Make ddtrace resolve to DIFFERENT ids so we can prove which backend won. + monkeypatch.setattr(obs_ids_mod, "_ddtrace_ids", lambda: ("d" * 32, "e" * 16)) + + # prefer_otel (Temporal path): OTel wins even though mode is dd_only. + assert obs_correlation(prefer_otel=True) == {"obs_trace_id": _TRACE_HEX, "obs_span_id": _SPAN_HEX} + # Default (in-process path): still honors mode -> ddtrace. + assert obs_correlation() == {"obs_trace_id": "d" * 32, "obs_span_id": "e" * 16} + + +def test_tag_ambient_prefer_otel_falls_back_to_ddtrace(monkeypatch: pytest.MonkeyPatch) -> None: + """When no valid OTel span is active, prefer_otel falls back to ddtrace.""" + monkeypatch.setenv("SGP_OBS_MODE", "dd_only") + + # No valid OTel span active. + invalid = _RecordingOtelSpan(otel_trace.INVALID_SPAN_CONTEXT) + monkeypatch.setattr(otel_trace, "get_current_span", lambda *a, **k: invalid) + + tagged: dict[str, Any] = {} + + class _FakeDDSpan: + def set_tag(self, k: str, v: Any) -> None: + tagged[k] = v + + class _FakeDDTracer: + def current_span(self) -> _FakeDDSpan: + return _FakeDDSpan() + + # obs_span imports `from ddtrace.trace import tracer` lazily; inject a stub module. + import sys + import types + + ddtrace_trace = types.ModuleType("ddtrace.trace") + ddtrace_trace.tracer = _FakeDDTracer() # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "ddtrace.trace", ddtrace_trace) + + tag_ambient_obs_span(business_span_id="bs", business_trace_id="bt", prefer_otel=True) + + # OTel was invalid -> fell back to ddtrace, which got the reverse tag. + assert tagged["agentex.business_span_id"] == "bs" + assert tagged["agentex.business_trace_id"] == "bt" + # The invalid OTel span was NOT tagged. + assert invalid.attributes == {} From 2a9576b58965e90a70ea200c972b3afdfb2ec742 Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Wed, 5 Aug 2026 23:48:49 -0700 Subject: [PATCH 11/12] test(tracing): update obs_span invalid-context test to the None-fallback contract The comment-#2 fix changed open_obs_span so an invalid wrapper context returns None (caller falls back to ambient obs_correlation()) instead of a handle with empty correlation. The pre-existing tests/lib/core/tracing/test_obs_span.py still asserted the old contract and failed CI. Update it to assert None + that the attached context is detached and the no-op span ended (no leak). Co-Authored-By: Claude Opus 4.8 --- tests/lib/core/tracing/test_obs_span.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/tests/lib/core/tracing/test_obs_span.py b/tests/lib/core/tracing/test_obs_span.py index 9a4ff69e4..a7f40a511 100644 --- a/tests/lib/core/tracing/test_obs_span.py +++ b/tests/lib/core/tracing/test_obs_span.py @@ -150,19 +150,29 @@ def test_lgtm_opens_named_span_and_reads_its_ids(self, monkeypatch): "agentex.business_trace_id": "btrace-1", } - def test_invalid_span_context_yields_empty_correlation(self, monkeypatch): + def test_invalid_span_context_returns_none_for_fallback(self, monkeypatch): + """Invalid wrapper context (proxy NonRecordingSpan / no TracerProvider): + open_obs_span returns None so the caller falls back to the ambient + obs_correlation() instead of taking an empty-correlation handle (which + would suppress the fallback and strip obs_* ids). It also detaches the + context it attached and ends the no-op span, so nothing leaks.""" monkeypatch.setenv("SGP_OBS_MODE", "lgtm") - _install_fake_otel(monkeypatch) + record = _install_fake_otel(monkeypatch) + + made: dict = {} def start_span(name): span = _FakeOtelSpan(name, 0, 0) span._ctx = _FakeSpanContext(0, 0, is_valid=False) + made["span"] = span return span sys.modules["opentelemetry"].trace.get_tracer = lambda _n: types.SimpleNamespace(start_span=start_span) handle = obs_span.open_obs_span("step") - assert handle is not None - assert handle.correlation == {} + assert handle is None + # cleaned up: the attached context was detached and the no-op span ended + assert len(record["detached"]) == 1 + assert made["span"].ended is True def test_close_detaches_and_ends(self, monkeypatch): monkeypatch.setenv("SGP_OBS_MODE", "lgtm") From a4f444ed247df68c3f67b8bd55325859886b3d9a Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Thu, 6 Aug 2026 14:26:26 -0700 Subject: [PATCH 12/12] fix(tracing): fail-open across the whole ACP-dispatch span setup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _acp_dispatch_span only guarded the opentelemetry import; get_tracer() and entering start_as_current_span (which runs the sampler and every SpanProcessor.on_start — the SDK guards neither) ran unprotected, so a broken provider or a custom sampler/processor that raises would fail the task/create / event/send dispatch itself, against the fail-open principle used elsewhere. Guard the full setup (import + get_tracer + span __enter__); on any failure run the dispatch untraced. Keep the dispatch body (yield) outside the guard so its exceptions still propagate, and guard __exit__ (passing exc info so the span reflects a failed dispatch) so closing can't mask the outcome. Co-Authored-By: Claude Opus 4.8 --- .../services/temporal_task_service.py | 37 ++++++++++++++----- 1 file changed, 28 insertions(+), 9 deletions(-) diff --git a/src/agentex/lib/core/temporal/services/temporal_task_service.py b/src/agentex/lib/core/temporal/services/temporal_task_service.py index d12f20eb8..20eb9d56e 100644 --- a/src/agentex/lib/core/temporal/services/temporal_task_service.py +++ b/src/agentex/lib/core/temporal/services/temporal_task_service.py @@ -1,5 +1,6 @@ from __future__ import annotations +import sys from typing import Any from datetime import timedelta from contextlib import contextmanager @@ -30,20 +31,38 @@ def _acp_dispatch_span(name: str, task_id: str | None = None) -> Iterator[None]: Opening a span here gives the interceptor something to inject. It becomes a child of the ingress request span when one is active (front-of-request - propagation), or a fresh per-turn root otherwise. Fail-open: never raises if - OpenTelemetry isn't importable. + propagation), or a fresh per-turn root otherwise. + + Fail-open across the WHOLE obs setup, not just the import: ``get_tracer`` and + entering ``start_as_current_span`` run the sampler and every + ``SpanProcessor.on_start`` (the SDK does not guard those), so a broken + provider or a custom sampler/processor that raises would otherwise fail the + dispatch itself. If any of it fails we run the dispatch untraced. The dispatch + body (the ``yield``) is OUTSIDE the guard so its exceptions still propagate. """ + span_cm = None try: from opentelemetry import trace as _otel_trace + + tracer = _otel_trace.get_tracer("agentex.acp") + # task_id goes on an attribute, NOT in the span name: a per-task span name is + # high-cardinality and breaks span-name aggregation in Tempo. + attributes = {"agentex.task_id": task_id} if task_id else None + span_cm = tracer.start_as_current_span(name, kind=_otel_trace.SpanKind.PRODUCER, attributes=attributes) + span_cm.__enter__() except Exception: # pragma: no cover - obs must never break a dispatch + span_cm = None + + try: yield - return - tracer = _otel_trace.get_tracer("agentex.acp") - # task_id goes on an attribute, NOT in the span name: a per-task span name is - # high-cardinality and breaks span-name aggregation in Tempo. - attributes = {"agentex.task_id": task_id} if task_id else None - with tracer.start_as_current_span(name, kind=_otel_trace.SpanKind.PRODUCER, attributes=attributes): - yield + finally: + if span_cm is not None: + # Pass exc info so the span reflects a failed dispatch; guard __exit__ + # so closing the span can never mask the dispatch outcome. + try: + span_cm.__exit__(*sys.exc_info()) + except Exception: # pragma: no cover - best-effort close + pass class TemporalTaskService: