[WIP] feat(tracing): correlate business spans with active observability trace (#465)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: NiteshDhanpal

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

@@ -0,0 +1,83 @@
+"""Correlate adk business spans with the active observability trace.
+
+The adk business ``trace_id`` is the agent **task id** (run-level: it spans the
+whole agent run across many requests -- task/create, then each message/send turn),
+so we must NOT overwrite it with a per-request observability trace_id. Doing so
+would collapse the run-level grouping.
+
+Instead, each business span is *tagged* with the active observability
+trace_id/span_id (this is the OpenTelemetry "span link" pattern -- correlate
+across trace granularities rather than merging them). You can then pivot from a
+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:
+  - unset / "dd_only": ddtrace context (current stack)
+  - "dual":            OTel/LGTM preferred, ddtrace fallback
+  - "lgtm":            OTel/LGTM only
+
+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
+from typing import Dict, Optional, Tuple
+
+__all__ = ("get_obs_mode", "obs_correlation")
+
+DD_ONLY = "dd_only"
+DUAL = "dual"
+LGTM = "lgtm"
+_DEFAULT_MODE = DD_ONLY
+_VALID_MODES = (DD_ONLY, DUAL, LGTM)
+
+
+def get_obs_mode() -> str:
+    """Unset/empty/unrecognized -> ``dd_only`` (current behavior)."""
+    raw = os.getenv("SGP_OBS_MODE")
+    if not raw:
+        return _DEFAULT_MODE
+    mode = raw.strip().lower()
+    return mode if mode in _VALID_MODES else _DEFAULT_MODE
+
+
+def _lgtm_ids() -> Optional[Tuple[str, str]]:
+    try:
+        from opentelemetry import trace
+    except ImportError:
+        return None
+    ctx = trace.get_current_span().get_span_context()
+    if ctx and ctx.is_valid:
+        return format(ctx.trace_id, "032x"), format(ctx.span_id, "016x")
+    return None
+
+
+def _ddtrace_ids() -> Optional[Tuple[str, str]]:
+    try:
+        from ddtrace import tracer
+    except ImportError:
+        return None
+    ctx = tracer.current_trace_context()
+    if ctx and ctx.trace_id:
+        return format(ctx.trace_id, "032x"), format(ctx.span_id or 0, "016x")
+    return None
+
+
+def obs_correlation() -> Dict[str, str]:
+    """Return ``{"obs.trace_id": ..., "obs.span_id": ...}`` for the active
+    observability context, or ``{}`` if none is active.
+
+    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()
+
+    if not ids:
+        return {}
+    return {"obs.trace_id": ids[0], "obs.span_id": ids[1]}

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

@@ -11,6 +11,7 @@
 from agentex.types.span import Span
 from agentex.lib.utils.logging import make_logger
 from agentex.lib.utils.model_utils import recursive_model_dump
+from agentex.lib.core.tracing.obs_ids import obs_correlation
 from agentex.lib.core.tracing.span_error import set_span_error
 from agentex.lib.core.tracing.span_queue import (
     SpanEventType,
@@ -79,6 +80,12 @@ 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()
+        if obs:
+            serialized_data = {**(serialized_data or {}), **obs}
         id = str(uuid.uuid4())
 
         span = Span(
@@ -229,6 +236,12 @@ 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()
+        if obs:
+            serialized_data = {**(serialized_data or {}), **obs}
         id = str(uuid.uuid4())
 
         span = Span(