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 5f6c0c381..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,7 +1,10 @@ from __future__ import annotations +import sys 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 @@ -13,6 +16,55 @@ from agentex.lib.core.clients.temporal.temporal_client import TemporalClient +@contextmanager +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 + 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 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 + 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: """ Submits Agent agent_tasks to the async runtime for execution. @@ -26,7 +78,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 +88,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("acp.task_create", task_id=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 +111,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("acp.event_send", task_id=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. diff --git a/src/agentex/lib/core/tracing/obs_ids.py b/src/agentex/lib/core/tracing/obs_ids.py index 99c6b2555..45fada783 100644 --- a/src/agentex/lib/core/tracing/obs_ids.py +++ b/src/agentex/lib/core/tracing/obs_ids.py @@ -11,14 +11,20 @@ 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. """ + from __future__ import annotations import os @@ -27,10 +33,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: @@ -64,20 +69,31 @@ def _ddtrace_ids() -> Optional[Tuple[str, str]]: return None -def obs_correlation() -> Dict[str, str]: - """Return ``{"obs.trace_id": ..., "obs.span_id": ...}`` for the active +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. + 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'``). + + ``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. """ - 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: + 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 {} 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..385507269 --- /dev/null +++ b/src/agentex/lib/core/tracing/obs_span.py @@ -0,0 +1,283 @@ +"""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 Dict, Callable, Optional + +from agentex.lib.core.tracing.obs_ids import LGTM, get_obs_mode + +__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" + +# 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[[Optional[Dict[str, str]]], None], + ): + 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.""" + 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 trace, context + 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() + 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: + 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: + 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. + ctx = tracer.current_trace_context() + if ctx is None: + return None + # 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 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) + + 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 + + +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 _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. + + 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. + + ``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: + _tag_otel_ambient(business_span_id, business_trace_id) + else: + _tag_ddtrace_ambient(business_span_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, +) -> 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(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 c3ec91bc3..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 @@ -12,7 +13,13 @@ 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.obs_span import ( + 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 ( SpanEventType, AsyncSpanQueue, @@ -25,6 +32,145 @@ 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. +# +# 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: + """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 + + +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: """ @@ -49,6 +195,9 @@ def __init__( self.processors = processors self.client = client self.trace_id = trace_id + # 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, @@ -80,13 +229,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() + # 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()) + obs_handle, obs = _begin_obs(name, id, self.trace_id) if obs: serialized_data = {**(serialized_data or {}), **obs} - id = str(uuid.uuid4()) span = Span( id=id, @@ -98,9 +246,11 @@ def start_span( data=serialized_data, task_id=task_id, ) + if obs_handle is not None: + _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 @@ -120,12 +270,16 @@ def end_span( if span.end_time is None: span.end_time = datetime.now(UTC) + # 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(_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 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 @@ -206,6 +360,9 @@ def __init__( self.client = client self.trace_id = trace_id self._span_queue = span_queue or get_default_span_queue() + # 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, @@ -236,13 +393,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() + # 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()) + obs_handle, obs = _begin_obs(name, id, self.trace_id) if obs: serialized_data = {**(serialized_data or {}), **obs} - id = str(uuid.uuid4()) span = Span( id=id, @@ -254,9 +410,21 @@ async def start_span( data=serialized_data, task_id=task_id, ) + if obs_handle is not None: + _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 @@ -276,12 +444,22 @@ async def end_span( if span.end_time is None: span.end_time = datetime.now(UTC) + # 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(_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 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/lib/core/tracing/test_obs_ids.py b/tests/lib/core/tracing/test_obs_ids.py new file mode 100644 index 000000000..5cdeb81b8 --- /dev/null +++ b/tests/lib/core/tracing/test_obs_ids.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +import sys +import types +from typing import Any + +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: 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) + + 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 + + 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: Any = types.ModuleType("opentelemetry") + fake_otel.trace = fake_trace_mod + monkeypatch.setitem(sys.modules, "opentelemetry", fake_otel) + + 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: 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) + + 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..a7f40a511 --- /dev/null +++ b/tests/lib/core/tracing/test_obs_span.py @@ -0,0 +1,516 @@ +from __future__ import annotations + +import sys +import types +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from agentex.lib.core.tracing import trace as trace_module, obs_span +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. +# --------------------------------------------------------------------------- # +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 _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 + + def end(self): + self.ended = True + + +def _install_fake_otel(monkeypatch, *, trace_id=0xABC, span_id=0xFF): + record: dict[str, Any] = {"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}, + Status=_FakeStatus, + StatusCode=_FakeStatusCode, + ) + fake_context = types.SimpleNamespace( + attach=lambda ctx: record["attached"].append(ctx) or object(), + detach=lambda token: record["detached"].append(token), + ) + fake_otel: Any = 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.error = 0 + 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[str, Any] = {"span": None, "started": []} + ctx_obj = object() if active else None + record["ctx"] = ctx_obj + + def start_span(name, child_of=None, activate=False): + span = _FakeDDSpan(name, trace_id, span_id) + record["span"] = span + record["started"].append({"name": name, "child_of": child_of, "activate": activate}) + return span + + tracer = types.SimpleNamespace( + current_trace_context=lambda: ctx_obj, + start_span=start_span, + ) + 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) + 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_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") + 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 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") + 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 + + 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) +# --------------------------------------------------------------------------- # +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" + 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:], + } + # 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 + + 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 +# --------------------------------------------------------------------------- # +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 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 + 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 == { + "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_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") + 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 isinstance(span.data, dict) + 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_module._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) + 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 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 + + +# --------------------------------------------------------------------------- # +# 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 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): + """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}, + Status=_FakeStatus, + StatusCode=_FakeStatusCode, + ) + fake_context = types.SimpleNamespace( + attach=lambda ctx: object(), + detach=lambda token: None, + ) + fake_otel: Any = 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} 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 == {}