From 8bbfd9c58fb6eb08db0a876c512f1118449add4b Mon Sep 17 00:00:00 2001 From: "stainless-app[bot]" <142633134+stainless-app[bot]@users.noreply.github.com> Date: Tue, 4 Aug 2026 00:28:11 +0000 Subject: [PATCH 01/12] codegen metadata --- .stats.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.stats.yml b/.stats.yml index 04c8e43b9..25881e3c6 100644 --- a/.stats.yml +++ b/.stats.yml @@ -1,4 +1,4 @@ configured_endpoints: 75 -openapi_spec_url: https://storage.googleapis.com/stainless-sdk-openapi-specs/sgp/agentex-sdk-644a4ec06aa1f055c614cbef3379684819a4edd84eeb20d2fb29ae01663622a3.yml -openapi_spec_hash: a6a4dc0c09691ac9783bf38e9653a464 +openapi_spec_url: https://storage.googleapis.com/stainless-sdk-openapi-specs/sgp/agentex-sdk-330ce4f0d8feed6caeb73d6b12277cfd89f6ad85535b8c8a6f509743b0b6f8cb.yml +openapi_spec_hash: ed6b33682c511df6de538714c0864aa3 config_hash: 593e89b291976a5e84e4c3c3f8324354 From da7ea1558683da05a3f9ecb119b91bf873437be1 Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Tue, 4 Aug 2026 16:02:15 -0700 Subject: [PATCH 02/12] feat(tracing): propagate OTel trace context across Temporal boundaries (#485) Co-authored-by: Claude Opus 4.8 --- .../lib/core/clients/temporal/utils.py | 5 ++ .../lib/core/temporal/workers/worker.py | 8 +- src/agentex/lib/core/tracing/temporal.py | 73 +++++++++++++++++++ .../core/tracing/test_temporal_interceptor.py | 40 ++++++++++ 4 files changed, 125 insertions(+), 1 deletion(-) create mode 100644 src/agentex/lib/core/tracing/temporal.py create mode 100644 tests/lib/core/tracing/test_temporal_interceptor.py diff --git a/src/agentex/lib/core/clients/temporal/utils.py b/src/agentex/lib/core/clients/temporal/utils.py index 95319720a..15b08cec6 100644 --- a/src/agentex/lib/core/clients/temporal/utils.py +++ b/src/agentex/lib/core/clients/temporal/utils.py @@ -9,6 +9,8 @@ from temporalio.converter import PayloadCodec, DataConverter from temporalio.contrib.pydantic import pydantic_data_converter +from agentex.lib.core.tracing.temporal import temporal_tracing_interceptors + # class DateTimeJSONEncoder(AdvancedJSONEncoder): # def default(self, o: Any) -> Any: # if isinstance(o, datetime.datetime): @@ -136,6 +138,9 @@ async def get_temporal_client( connect_kwargs: dict[str, Any] = { "target_host": temporal_address, "plugins": plugins, + # Propagate OTel trace context on outbound start_workflow / execute_activity + # (enabled by default; AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED=false to disable). + "interceptors": temporal_tracing_interceptors(), } if data_converter is not None: diff --git a/src/agentex/lib/core/temporal/workers/worker.py b/src/agentex/lib/core/temporal/workers/worker.py index 2b4958b1f..0cfe01185 100644 --- a/src/agentex/lib/core/temporal/workers/worker.py +++ b/src/agentex/lib/core/temporal/workers/worker.py @@ -29,6 +29,7 @@ from agentex.lib.utils.logging import make_logger from agentex.lib.utils.registration import register_agent +from agentex.lib.core.tracing.temporal import temporal_tracing_interceptors from agentex.lib.environment_variables import EnvironmentVariables from agentex.lib.core.compat.version_guard import assert_backend_compatible @@ -126,6 +127,9 @@ async def get_temporal_client( connect_kwargs: dict[str, Any] = { "target_host": temporal_address, "plugins": plugins, + # Propagate OTel trace context on outbound start_workflow / execute_activity + # (enabled by default; AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED=false to disable). + "interceptors": temporal_tracing_interceptors(), } if data_converter is not None: @@ -229,7 +233,9 @@ async def run( max_concurrent_activities=self.max_concurrent_activities, build_id=str(uuid.uuid4()), debug_mode=debug_enabled, # Disable deadlock detection in debug mode - interceptors=self.interceptors, # Pass interceptors to Worker + # Tracing interceptor OUTERMOST so business interceptors (and the spans + # they create) nest under the propagated workflow/activity span. + interceptors=[*temporal_tracing_interceptors(), *self.interceptors], ) logger.info(f"Starting workers for task queue: {self.task_queue}") diff --git a/src/agentex/lib/core/tracing/temporal.py b/src/agentex/lib/core/tracing/temporal.py new file mode 100644 index 000000000..484abc26b --- /dev/null +++ b/src/agentex/lib/core/tracing/temporal.py @@ -0,0 +1,73 @@ +"""OpenTelemetry trace-context propagation across Temporal boundaries. + +Temporal serializes ``start_workflow`` / ``execute_activity`` across (potentially +cross-process) boundaries, and does NOT carry the active W3C ``traceparent`` by +default. So any span created inside a workflow or activity becomes a **new +detached root** -- the trace shatters at every Temporal hop. + +This bites agentex directly: ``adk.tracing.span`` runs span creation as a +Temporal activity when ``in_temporal_workflow()`` is true, so without propagation +those business spans detach from the turn's obs trace. + +Wiring temporalio's first-party ``TracingInterceptor`` onto the Temporal client +and worker injects the active span context into Temporal headers on the caller +side and extracts + continues it on the workflow/activity side, using the global +OpenTelemetry propagator -- so ``client -> workflow -> activity`` is one trace. + +Enabled by DEFAULT. Set ``AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED=false`` +(also accepts ``0`` / ``no`` / ``off``) to turn it off. It also degrades to a +no-op -- and never raises -- if temporalio's OpenTelemetry contrib isn't +importable, so enabling it by default can't break a worker. +""" + +from __future__ import annotations + +import os +from typing import Any + +from agentex.lib.utils.logging import make_logger + +logger = make_logger(__name__) + +_ENABLE_ENV = "AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED" +_FALSEY = {"0", "false", "no", "off"} + + +def temporal_trace_interceptor_enabled() -> bool: + """Whether the Temporal OTel trace interceptor should be installed. + + Defaults to True; disabled only when ``AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED`` + is set to a falsy value (``0`` / ``false`` / ``no`` / ``off``).""" + return os.environ.get(_ENABLE_ENV, "true").strip().lower() not in _FALSEY + + +def temporal_tracing_interceptors() -> list[Any]: + """Interceptors that propagate OpenTelemetry trace context across Temporal. + + Returns ``[TracingInterceptor()]`` (enabled by default) so callers can splat + it into a client's / worker's ``interceptors=`` list. Returns ``[]`` when + disabled via env, or when temporalio's OpenTelemetry contrib is not + importable. Never raises -- observability wiring must not break a worker. + + ``TracingInterceptor`` implements both the client and worker interceptor + interfaces, so the same call is used on both sides: + - on the **client**, it injects context on outbound ``start_workflow`` / + ``execute_activity`` calls; + - on the **worker**, it extracts context and roots the workflow / activity + execution spans under it. + """ + if not temporal_trace_interceptor_enabled(): + logger.info("Temporal OTel trace interceptor disabled via %s", _ENABLE_ENV) + return [] + try: + from temporalio.contrib.opentelemetry import TracingInterceptor + + # Construct inside the try so a constructor failure (not just a missing + # contrib) also falls back to a no-op instead of aborting worker startup. + return [TracingInterceptor()] + except Exception as exc: # contrib unavailable OR constructor failure -> no-op, never raise + logger.warning( + "Temporal OTel trace interceptor unavailable (%s); traces will not propagate across Temporal boundaries.", + exc, + ) + return [] diff --git a/tests/lib/core/tracing/test_temporal_interceptor.py b/tests/lib/core/tracing/test_temporal_interceptor.py new file mode 100644 index 000000000..83c0c3681 --- /dev/null +++ b/tests/lib/core/tracing/test_temporal_interceptor.py @@ -0,0 +1,40 @@ +"""Unit tests for the Temporal OTel trace-interceptor wiring. + +Verifies the interceptor is on by default, the opt-out env flag, and the safe +no-op fallback when temporalio's OpenTelemetry contrib isn't importable. +""" + +import sys + +import pytest + +from agentex.lib.core.tracing import temporal as temporal_tracing + + +class TestTemporalTraceInterceptor: + def test_enabled_by_default(self, monkeypatch): + monkeypatch.delenv("AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED", raising=False) + assert temporal_tracing.temporal_trace_interceptor_enabled() is True + + interceptors = temporal_tracing.temporal_tracing_interceptors() + assert len(interceptors) == 1 + # temporalio's first-party OTel interceptor + assert type(interceptors[0]).__name__ == "TracingInterceptor" + + @pytest.mark.parametrize("value", ["false", "0", "no", "off", "FALSE", "Off"]) + def test_disabled_via_env(self, monkeypatch, value): + monkeypatch.setenv("AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED", value) + assert temporal_tracing.temporal_trace_interceptor_enabled() is False + assert temporal_tracing.temporal_tracing_interceptors() == [] + + @pytest.mark.parametrize("value", ["true", "1", "yes", "TRUE", "anything"]) + def test_enabled_for_non_falsy_values(self, monkeypatch, value): + monkeypatch.setenv("AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED", value) + assert temporal_tracing.temporal_trace_interceptor_enabled() is True + + def test_no_op_when_contrib_unimportable(self, monkeypatch): + # Enabled, but temporalio's OTel contrib not importable -> [] (never raises), + # so default-on can't break a worker that lacks the contrib. + monkeypatch.delenv("AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED", raising=False) + monkeypatch.setitem(sys.modules, "temporalio.contrib.opentelemetry", None) + assert temporal_tracing.temporal_tracing_interceptors() == [] From 72732b7c07700df840a2308424154f11a30e39f2 Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Thu, 6 Aug 2026 14:32:41 -0700 Subject: [PATCH 03/12] feat(tracing): correlate business spans with obs via dedicated wrapper span (#484) Co-authored-by: Claude Opus 4.8 --- src/agentex/lib/adk/_modules/tracing.py | 19 + .../services/temporal_task_service.py | 103 +++- src/agentex/lib/core/tracing/obs_ids.py | 44 +- src/agentex/lib/core/tracing/obs_span.py | 283 ++++++++++ src/agentex/lib/core/tracing/trace.py | 208 ++++++- tests/lib/core/tracing/test_obs_ids.py | 126 +++++ tests/lib/core/tracing/test_obs_span.py | 516 ++++++++++++++++++ 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 +++++ 11 files changed, 1727 insertions(+), 56 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 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 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 == {} From 7b94c34c542c702e3897a85f4477bee6e37343cf Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Thu, 6 Aug 2026 23:43:41 -0700 Subject: [PATCH 04/12] feat(tracing): per-step obs wrappers inside business Temporal activities (1:1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously _begin_obs skipped the obs wrapper for ANY Temporal activity (Option A) and only stamped the ambient RunActivity span, so all business spans in a turn collapsed onto ONE obs span (52:1). But inside a *business* activity, start_span and end_span run in the SAME process, so a wrapper is safe there. Option A is only required for the SDK's own dispatched START_SPAN/END_SPAN activities (the in_temporal_workflow path), where start and end are separate activities on possibly different workers. Discriminate on activity type: _in_tracing_dispatch_activity() is true only for the "start-span"/"end-span" activities. For everything else (sync, or a business activity) open a real per-step wrapper — it nests under the interceptor's ambient RunActivity span and closes in-process, giving each business span its own obs span (1:1), matching the sync path. The bounded _OBS_HANDLES registry backstops any mis-discrimination. --- src/agentex/lib/core/tracing/trace.py | 70 ++++++++++++--------------- tests/test_temporal_obs_backend.py | 35 +++++++++++++- 2 files changed, 64 insertions(+), 41 deletions(-) diff --git a/src/agentex/lib/core/tracing/trace.py b/src/agentex/lib/core/tracing/trace.py index d3decdb9b..447ed7d9b 100644 --- a/src/agentex/lib/core/tracing/trace.py +++ b/src/agentex/lib/core/tracing/trace.py @@ -106,39 +106,25 @@ def _run_on_span_end(processor: SyncTracingProcessor, span: Span) -> None: ) -def _in_temporal_activity() -> bool: - """True when executing inside a Temporal activity. - - On the Temporal path ``start_span`` and ``end_span`` run as SEPARATE - activities (START_SPAN / END_SPAN) that Temporal can route to DIFFERENT - worker processes. A wrapper obs span opened in the START_SPAN activity could - therefore never be closed by END_SPAN -- its handle lives in another - process's ``_OBS_HANDLES`` -- so it would leak (unbounded, OOM risk) and its - persisted ``obs_span_id`` would dangle (the span is never .end()ed, so never - exported to Tempo). - - So inside an activity we do NOT open our own wrapper. We lean on the span the - Temporal OTel ``TracingInterceptor`` (see ``core/tracing/temporal.py`` + - scale-agentex-python#485) already made active for this activity -- which is - rooted under the turn's propagated trace -- and merely stamp the reverse tag - onto it (``tag_ambient_obs_span``). That keeps trace-level correlation with - no cross-process handle to leak. - - Never raises; returns False when temporalio isn't importable. - - TODO(obs-followup): this intentionally drops the *named per-step* wrapper on - the Temporal path (obs_span_id becomes the ambient activity span, not a - step-named span) and does NOT add TurnTrace RETRY/ASYNC roll-up -- retried - turns still surface as N unlinked spans. Follow-up diff should (a) optionally - materialize a self-contained named wrapper inside a single activity using the - span's own start/end timestamps, and (b) build the TurnTrace roll-up. - Test-later: on a multi-replica worker fleet, assert _OBS_HANDLES stays - bounded (no leak / OOM) and that obs_trace_id resolves to the turn trace. - """ +def _in_tracing_dispatch_activity() -> bool: + """True only when running inside the SDK's OWN dispatched START_SPAN / END_SPAN + activity (the ``in_temporal_workflow()`` path, where a workflow runs span start + and end as SEPARATE activities that Temporal can route to different workers). + + That is the one case a per-step obs wrapper can't work: the wrapper opened in + the START_SPAN activity could never be closed by the END_SPAN activity. A span + created directly inside a *business* activity (an agent turn's own + ``adk.tracing.span``) runs start AND end in the same activity process, so a + wrapper there is safe -- it nests under the interceptor's ambient RunActivity + span and closes in-process. The tracing dispatch activities are named + ``start-span`` / ``end-span`` (``TracingActivityName``). Never raises; False + when temporalio isn't importable or we're not in an activity.""" try: from temporalio import activity - return activity.in_activity() + if not activity.in_activity(): + return False + return activity.info().activity_type in ("start-span", "end-span") except Exception: return False @@ -148,23 +134,27 @@ def _begin_obs( span_id: str, trace_id: str | None, ) -> tuple[ObsSpanHandle | None, dict[str, str]]: - """Open the obs wrapper for a business span (or, inside a Temporal activity, - tag the ambient interceptor span) and return ``(handle, correlation)``. + """Open the obs wrapper for a business span and return ``(handle, correlation)``. Shared by ``Trace.start_span`` and ``AsyncTrace.start_span`` so the two paths can't drift. The wrapper is named for the step so ``obs_span_id`` is stable/meaningful (not an arbitrary innermost httpx span), and it carries the reverse tag (business span/trace id) for the obs -> business pivot. - Temporal path: we do NOT open our own wrapper -- start_span / end_span run as - separate activities on possibly different workers, so the handle could never - be closed. Instead we tag the span the temporalio OTel ``TracingInterceptor`` - already made active. That span is OTel REGARDLESS of ``SGP_OBS_MODE``, so we - pass ``prefer_otel=True`` to both the tag and the correlation read -- otherwise - the default ``dd_only`` mode would tag/read an unrelated ddtrace span and the - ids would point at the wrong trace. See ``_in_temporal_activity``. + We open a real per-step wrapper on the sync path AND inside a *business* + Temporal activity -- there the wrapper nests under the interceptor's ambient + RunActivity span and start/end run in-process, so it closes cleanly and each + business step gets its own obs span (1:1), just like sync. + + The ONE exception is the SDK's own dispatched START_SPAN / END_SPAN activity + (a workflow calling ``adk.tracing`` -- see ``_in_tracing_dispatch_activity``): + there start and end are separate activities on possibly different workers, so + a wrapper could never be closed. We fall back to tagging the ambient + interceptor span instead, with ``prefer_otel=True`` (the interceptor span is + OTel regardless of ``SGP_OBS_MODE``, so a plain ``dd_only`` read would + otherwise point at an unrelated ddtrace span). """ - if _in_temporal_activity(): + if _in_tracing_dispatch_activity(): tag_ambient_obs_span(business_span_id=span_id, business_trace_id=trace_id, prefer_otel=True) return None, obs_correlation(prefer_otel=True) handle = open_obs_span(name, business_span_id=span_id, business_trace_id=trace_id) diff --git a/tests/test_temporal_obs_backend.py b/tests/test_temporal_obs_backend.py index 34daf1d11..1ca1f5959 100644 --- a/tests/test_temporal_obs_backend.py +++ b/tests/test_temporal_obs_backend.py @@ -67,8 +67,10 @@ def _activate_otel_span(monkeypatch: pytest.MonkeyPatch) -> _RecordingOtelSpan: def test_temporal_path_tags_and_reads_otel_in_dd_only(monkeypatch: pytest.MonkeyPatch) -> None: # Default/dd_only mode is exactly where the old code went to ddtrace. + # Option A (tag the ambient interceptor span, no wrapper) now applies only + # inside the SDK's dispatched START_SPAN/END_SPAN activity, not any activity. monkeypatch.setenv("SGP_OBS_MODE", "dd_only") - monkeypatch.setattr(trace_mod, "_in_temporal_activity", lambda: True) + monkeypatch.setattr(trace_mod, "_in_tracing_dispatch_activity", lambda: True) activity_span = _activate_otel_span(monkeypatch) trace_obj = Trace(processors=[], client=cast(Any, object()), trace_id="trace-1") @@ -132,3 +134,34 @@ def current_span(self) -> _FakeDDSpan: assert tagged["agentex.business_trace_id"] == "bt" # The invalid OTel span was NOT tagged. assert invalid.attributes == {} + + +class _FakeHandle: + def __init__(self, corr): + self.correlation = corr + + +def test_begin_obs_opens_wrapper_outside_dispatch_activity(monkeypatch: pytest.MonkeyPatch) -> None: + """Sync path or inside a business Temporal activity: open a per-step wrapper + (1:1), NOT Option A. Each business span gets its own obs span.""" + monkeypatch.setattr(trace_mod, "_in_tracing_dispatch_activity", lambda: False) + monkeypatch.setattr( + trace_mod, "open_obs_span", + lambda *a, **k: _FakeHandle({"obs_trace_id": "t1", "obs_span_id": "s1"}), + ) + handle, corr = trace_mod._begin_obs("mortgage.classify_intent", "bs", "bt") + assert handle is not None + assert corr == {"obs_trace_id": "t1", "obs_span_id": "s1"} + + +def test_begin_obs_tags_ambient_inside_dispatch_activity(monkeypatch: pytest.MonkeyPatch) -> None: + """Inside the dispatched START_SPAN/END_SPAN activity: no wrapper (would leak + across activities); tag the ambient interceptor span instead (Option A).""" + monkeypatch.setattr(trace_mod, "_in_tracing_dispatch_activity", lambda: True) + tagged: dict = {} + monkeypatch.setattr(trace_mod, "tag_ambient_obs_span", lambda **k: tagged.update(k)) + monkeypatch.setattr(trace_mod, "obs_correlation", lambda **k: {"obs_trace_id": "amb", "obs_span_id": "amb"}) + handle, corr = trace_mod._begin_obs("mortgage.advisor.turn", "bs", "bt") + assert handle is None + assert tagged.get("business_span_id") == "bs" and tagged.get("prefer_otel") is True + assert corr == {"obs_trace_id": "amb", "obs_span_id": "amb"} From 777a2e8c1a74ce45552730f56b0e3692732cbf61 Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Sun, 16 Aug 2026 23:11:57 -0700 Subject: [PATCH 05/12] feat(tracing): delegate obs edge to sgp-obs + wire init_tracing at ACP startup Replace the SDK's inline observability edge (#484/#485/#490/#491) with thin shims that delegate to the shared sgp-obs library: - obs_ids.py / obs_span.py / temporal.py / trace.py: delegate correlation, the per-step wrapper span, the begin-obs decision, and the Temporal interceptors to sgp_obs.traces backends + Correlator. - adk/pyproject.toml: add sgp-obs==0.2.0rc1; root pyproject: add the CodeArtifact `scale` index + explicit source for sgp-obs. Also wire sgp_obs.traces.init_tracing() into the BaseACPServer lifespan. Without it, get_tracer() resolves to the API-default ProxyTracerProvider: obs wrapper spans never record or export, so no span carrying the .business_trace_id reverse anchor reaches the collector. init_tracing adopts an app-installed provider if present, else installs one with the OTLP exporter (fail-open). Validated end-to-end in sgp-dev (rocket-mock + audit): agent registers, 20/20 business spans correlate (obs ids in span metadata), per-step wrapper spans record with the reverse anchor, and force_flush to the collector succeeds. Co-Authored-By: Claude Opus 4.8 --- adk/pyproject.toml | 4 + pyproject.toml | 13 + src/agentex/lib/core/tracing/obs_ids.py | 105 +++--- src/agentex/lib/core/tracing/obs_span.py | 321 +++++------------- src/agentex/lib/core/tracing/temporal.py | 73 ++-- src/agentex/lib/core/tracing/trace.py | 14 +- .../lib/sdk/fastacp/base/base_acp_server.py | 23 ++ uv.lock | 19 +- 8 files changed, 212 insertions(+), 360 deletions(-) diff --git a/adk/pyproject.toml b/adk/pyproject.toml index 3569a8e52..10e8a74e9 100644 --- a/adk/pyproject.toml +++ b/adk/pyproject.toml @@ -59,6 +59,10 @@ dependencies = [ "ddtrace>=3.13.0", "opentelemetry-api>=1.20.0", "opentelemetry-sdk>=1.20.0", + # SGP unified observability library: the obs edge (correlation, wrapper span, + # Temporal/ingress propagation) is delegated to this shared lib. rc pin while + # it stabilizes; flip to ==0.3.0 once that release publishes. + "sgp-obs==0.2.0rc1", "json_log_formatter>=1.1.1", ] diff --git a/pyproject.toml b/pyproject.toml index f73829573..fbc020bcb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,10 +54,23 @@ dev = [ # `src/agentex/lib/` ships only from the heavy (via adk/hatch_build.py). members = ["adk"] +[[tool.uv.index]] +# CodeArtifact scale-pypi, for internal resolution of sgp-obs. `explicit` so ONLY +# packages that name it via [tool.uv.sources] use it — every other dep still +# resolves from public PyPI. Auth via UV_INDEX_SCALE_USERNAME/PASSWORD (a +# CodeArtifact token). External consumers resolve sgp-obs from public PyPI once +# the dual-publish is live; this pin is a dev/CI resolution hint only. +name = "scale" +url = "https://scale-307185671274.d.codeartifact.us-west-2.amazonaws.com/pypi/scale-pypi/simple/" +explicit = true + [tool.uv.sources] # Dev-only: resolve the ADK's agentex-client dep to this root package. # Stripped from published wheels — the heavy wheel still pins the PyPI version. agentex-client = { workspace = true } +# Pull sgp-obs from the scale CodeArtifact index (a dev/CI hint; the published +# wheel just pins sgp-obs==0.2.0rc1 and lets the consumer's index resolve it). +sgp-obs = { index = "scale" } [tool.uv] managed = true diff --git a/src/agentex/lib/core/tracing/obs_ids.py b/src/agentex/lib/core/tracing/obs_ids.py index 45fada783..947bd36ba 100644 --- a/src/agentex/lib/core/tracing/obs_ids.py +++ b/src/agentex/lib/core/tracing/obs_ids.py @@ -1,45 +1,45 @@ """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: - - unset / "dd_only": ddtrace context (current stack) - - "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. +DELEGATES to ``sgp_obs`` (the shared SGP observability library). The id-reading +that used to live here (``_lgtm_ids`` / ``_ddtrace_ids``) now lives in +``sgp_obs.traces.backends``; this module keeps the SDK's public surface +(``get_obs_mode`` / ``obs_correlation`` + the ``DD_ONLY`` / ``LGTM`` constants) so +callers are unchanged. + +Each business span is *tagged* with the active observability trace_id/span_id +(the OpenTelemetry "span link" pattern) so you can pivot a persisted business +span to its per-turn Tempo/Datadog trace, while the business trace still groups +the whole run by task id. Never fabricates ids — no active context returns ``{}``. + +Source follows ``SGP_OBS_MODE``: unset/``dd_only`` -> ddtrace; ``lgtm`` -> OTel. +Note the SDK spells the OTel mode ``"lgtm"`` while ``sgp_obs.ObsMode`` spells it +``"otel"`` — :func:`obs_mode` maps between them. """ from __future__ import annotations import os -from typing import Dict, Tuple, Optional +from typing import Dict -__all__ = ("get_obs_mode", "obs_correlation") +from sgp_obs.traces import ObsMode +from sgp_obs.traces.backends.otel import OTelBackend +from sgp_obs.traces.backends.ddtrace import DDTraceBackend + +__all__ = ("get_obs_mode", "obs_correlation", "DD_ONLY", "LGTM") DD_ONLY = "dd_only" LGTM = "lgtm" _DEFAULT_MODE = DD_ONLY _VALID_MODES = (DD_ONLY, LGTM) +# Stateless singletons — the backends hold no per-call state. +_OTEL = OTelBackend() +_DDTRACE = DDTraceBackend() + def get_obs_mode() -> str: - """Unset/empty/unrecognized -> ``dd_only`` (current behavior).""" + """Unset/empty/unrecognized -> ``dd_only``. Returns the SDK's string form + (``"lgtm"`` / ``"dd_only"``).""" raw = os.getenv("SGP_OBS_MODE") if not raw: return _DEFAULT_MODE @@ -47,53 +47,26 @@ def get_obs_mode() -> str: 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.trace 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_mode() -> ObsMode: + """The SDK's ``SGP_OBS_MODE`` mapped to ``sgp_obs.ObsMode`` (``lgtm`` -> OTEL, + else DD_ONLY). Used to drive the sgp_obs Correlator/backends.""" + return ObsMode.OTEL if get_obs_mode() == LGTM else ObsMode.DD_ONLY 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. + """Return ``{"obs_trace_id": ..., "obs_span_id": ...}`` for the active obs + context, or ``{}`` if none. Delegates id-reading to the sgp_obs backends' + ``current_ids()``. - Never fabricates ids -- this is a correlation tag, not the span's id. + ``prefer_otel``: read OTel first (on the Temporal path the active span is the + temporalio OTel interceptor span regardless of ``SGP_OBS_MODE``). Never + fabricates ids — this is a correlation tag, not the span's id. """ try: if prefer_otel: - ids = _lgtm_ids() or _ddtrace_ids() + corr = _OTEL.current_ids() or _DDTRACE.current_ids() else: - ids = _lgtm_ids() if get_obs_mode() == LGTM else _ddtrace_ids() + corr = (_OTEL if get_obs_mode() == LGTM else _DDTRACE).current_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 corr.as_metadata() if corr is not None else {} diff --git a/src/agentex/lib/core/tracing/obs_span.py b/src/agentex/lib/core/tracing/obs_span.py index 385507269..55742b677 100644 --- a/src/agentex/lib/core/tracing/obs_span.py +++ b/src/agentex/lib/core/tracing/obs_span.py @@ -1,174 +1,66 @@ -"""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. +"""Dedicated per-business-span observability wrapper span — DELEGATED to sgp_obs. + +The span/id mechanics that used to live here (``_open_otel_span`` / +``_open_ddtrace_span`` / ``_tag_*_ambient`` / ``_hex_ids``) now live in +``sgp_obs.traces.backends`` behind the ``ObsBackend`` port, and the +wrapper-vs-ambient DECISION lives in ``sgp_obs.traces.Correlator``. This module +keeps the SDK's public surface (``ObsSpanHandle`` / ``open_obs_span`` / +``close_obs_span`` / ``tag_ambient_obs_span``) as thin adapters so callers and +tests are unchanged, and adds :func:`begin_obs` which ``trace.py`` uses to +delegate the whole decision in one call. + +Reverse-tag source is ``agentex`` (this is the agentex SDK). Never raises. """ from __future__ import annotations -from typing import Dict, Callable, Optional - -from agentex.lib.core.tracing.obs_ids import LGTM, get_obs_mode +from typing import Dict, Optional -__all__ = ("ObsSpanHandle", "open_obs_span", "close_obs_span", "tag_ambient_obs_span") +from sgp_obs.traces import SpanError, Correlator, BusinessRef, SpanRequest, BusinessSource, temporal as _sgp_temporal +from sgp_obs.traces.ports import ObsSpanHandle as _SgpHandle -# Instrumentation scope name so these wrapper spans are identifiable in Tempo/DD. -_TRACER_NAME = "agentex.business" +from agentex.lib.core.tracing.obs_ids import LGTM, _OTEL, _DDTRACE, obs_mode, get_obs_mode -# 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" +__all__ = ( + "ObsSpanHandle", + "open_obs_span", + "close_obs_span", + "tag_ambient_obs_span", + "begin_obs", +) -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 _to_span_error(error: Optional[Dict[str, str]]) -> Optional[SpanError]: + """SDK error dict ``{"type","message"}`` -> sgp_obs SpanError.""" + if not error: + return None + return SpanError(type=error.get("type"), message=error.get("message")) -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 _biz(span_id: Optional[str], trace_id: Optional[str]) -> BusinessRef: + # source=agentex: this SDK is the agentex business source; the reverse tag is + # therefore ``agentex.business_span_id`` / ``agentex.business_trace_id``. + return BusinessRef(trace_id=trace_id, span_id=span_id or "", source=BusinessSource.AGENTEX) -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 _backend(): + """The backend for the current SGP_OBS_MODE (lgtm -> OTel, else ddtrace).""" + return _OTEL if get_obs_mode() == LGTM else _DDTRACE - 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 +class ObsSpanHandle: + """Live handle for an open wrapper span. Adapts a sgp_obs ``ObsSpanHandle`` to + the SDK's shape: ``.correlation`` is the ``{obs_trace_id, obs_span_id}`` dict, + and ``close(error)`` takes the SDK's ``{"type","message"}`` error dict.""" + __slots__ = ("correlation", "_inner") -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 __init__(self, inner: _SgpHandle) -> None: + self._inner = inner + self.correlation: Dict[str, str] = inner.correlation.as_metadata() - 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 close(self, error: Optional[Dict[str, str]] = None) -> None: + self._inner.close(_to_span_error(error)) def open_obs_span( @@ -176,61 +68,14 @@ def open_obs_span( 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. - """ + """Open a step-named wrapper span in the active backend and make it active; + ``None`` when there's no live context (caller falls back to ambient ids). + Delegates to the sgp_obs backend. Never raises.""" 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) + inner = _backend().open_span(name, _biz(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 + return ObsSpanHandle(inner) if inner is not None else None def tag_ambient_obs_span( @@ -238,46 +83,52 @@ def tag_ambient_obs_span( 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.""" + """Stamp the reverse tag onto the CURRENTLY ACTIVE span without opening one. + ``prefer_otel`` tags OTel first (the Temporal interceptor span is OTel + regardless of SGP_OBS_MODE). Best-effort; never raises.""" + biz = _biz(business_span_id, business_trace_id) try: if prefer_otel: - if _tag_otel_ambient(business_span_id, business_trace_id): + if _OTEL.tag_ambient(biz): return - _tag_ddtrace_ambient(business_span_id, business_trace_id) + _DDTRACE.tag_ambient(biz) 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 + _backend().tag_ambient(biz) + except Exception: # pragma: no cover - best-effort 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``.""" +def close_obs_span(handle: Optional[ObsSpanHandle], error: Optional[Dict[str, str]] = None) -> None: + """Close the wrapper span, marking it errored when ``error`` is given. Safe on None.""" if handle is None: return try: handle.close(error) except Exception: # pragma: no cover - best-effort pass + + +def begin_obs( + name: str, + span_id: str, + trace_id: Optional[str], +) -> tuple[Optional[ObsSpanHandle], Dict[str, str]]: + """Delegate the ENTIRE wrapper-vs-ambient + backend decision to the sgp_obs + ``Correlator`` and return the SDK-shaped ``(handle | None, {obs_trace_id, + obs_span_id})``. Inside the dispatched start-span/end-span pair the Correlator + tags the ambient span (no wrapper); elsewhere it opens a per-step wrapper. + + A fresh Correlator per call keeps the SDK's per-call ``SGP_OBS_MODE`` reading + (the backends are stateless singletons; only the one-shot drift log resets). + """ + req = SpanRequest( + name=name, + business=_biz(span_id, trace_id), + in_activity=_sgp_temporal.in_activity(), + is_dispatch_boundary=_sgp_temporal.in_dispatch_boundary(), + ) + try: + inner, edge = Correlator(_OTEL, _DDTRACE, obs_mode()).begin(req) + except Exception: # pragma: no cover - obs must never break the business span + return None, {} + return (ObsSpanHandle(inner) if inner is not None else None), edge.as_metadata() diff --git a/src/agentex/lib/core/tracing/temporal.py b/src/agentex/lib/core/tracing/temporal.py index 484abc26b..32deca81a 100644 --- a/src/agentex/lib/core/tracing/temporal.py +++ b/src/agentex/lib/core/tracing/temporal.py @@ -1,23 +1,16 @@ -"""OpenTelemetry trace-context propagation across Temporal boundaries. - -Temporal serializes ``start_workflow`` / ``execute_activity`` across (potentially -cross-process) boundaries, and does NOT carry the active W3C ``traceparent`` by -default. So any span created inside a workflow or activity becomes a **new -detached root** -- the trace shatters at every Temporal hop. - -This bites agentex directly: ``adk.tracing.span`` runs span creation as a -Temporal activity when ``in_temporal_workflow()`` is true, so without propagation -those business spans detach from the turn's obs trace. - -Wiring temporalio's first-party ``TracingInterceptor`` onto the Temporal client -and worker injects the active span context into Temporal headers on the caller -side and extracts + continues it on the workflow/activity side, using the global -OpenTelemetry propagator -- so ``client -> workflow -> activity`` is one trace. - -Enabled by DEFAULT. Set ``AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED=false`` -(also accepts ``0`` / ``no`` / ``off``) to turn it off. It also degrades to a -no-op -- and never raises -- if temporalio's OpenTelemetry contrib isn't -importable, so enabling it by default can't break a worker. +"""OpenTelemetry trace-context propagation across Temporal boundaries — DELEGATED. + +The interceptor construction now lives in ``sgp_obs.traces.temporal``; this module +keeps the SDK's public surface (``temporal_trace_interceptor_enabled`` / +``temporal_tracing_interceptors``) and its +``AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED`` env toggle, delegating the actual +``TracingInterceptor`` to sgp_obs. + +Temporal serializes ``start_workflow`` / ``execute_activity`` across (possibly +cross-process) boundaries and does not carry the active W3C ``traceparent`` by +default, so a span made inside a workflow/activity would otherwise become a new +detached root. temporalio's first-party interceptor injects + continues the +context so ``client -> workflow -> activity`` is one trace. """ from __future__ import annotations @@ -25,6 +18,8 @@ import os from typing import Any +from sgp_obs.traces import temporal as _sgp_temporal + from agentex.lib.utils.logging import make_logger logger = make_logger(__name__) @@ -34,40 +29,22 @@ def temporal_trace_interceptor_enabled() -> bool: - """Whether the Temporal OTel trace interceptor should be installed. - - Defaults to True; disabled only when ``AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED`` - is set to a falsy value (``0`` / ``false`` / ``no`` / ``off``).""" + """Whether the Temporal OTel trace interceptor should be installed. Defaults + to True; disabled only when ``AGENTEX_TEMPORAL_TRACE_INTERCEPTOR_ENABLED`` is a + falsy value (``0`` / ``false`` / ``no`` / ``off``).""" return os.environ.get(_ENABLE_ENV, "true").strip().lower() not in _FALSEY def temporal_tracing_interceptors() -> list[Any]: - """Interceptors that propagate OpenTelemetry trace context across Temporal. - - Returns ``[TracingInterceptor()]`` (enabled by default) so callers can splat - it into a client's / worker's ``interceptors=`` list. Returns ``[]`` when - disabled via env, or when temporalio's OpenTelemetry contrib is not - importable. Never raises -- observability wiring must not break a worker. + """Interceptors that propagate OpenTelemetry context across Temporal. - ``TracingInterceptor`` implements both the client and worker interceptor - interfaces, so the same call is used on both sides: - - on the **client**, it injects context on outbound ``start_workflow`` / - ``execute_activity`` calls; - - on the **worker**, it extracts context and roots the workflow / activity - execution spans under it. + Returns ``[TracingInterceptor()]`` (from sgp_obs, which implements both the + client and worker interfaces) so callers splat it into a client's / worker's + ``interceptors=``. Returns ``[]`` when disabled via env, or when temporalio's + OpenTelemetry contrib isn't importable. Never raises — obs wiring must not + break a worker. """ if not temporal_trace_interceptor_enabled(): logger.info("Temporal OTel trace interceptor disabled via %s", _ENABLE_ENV) return [] - try: - from temporalio.contrib.opentelemetry import TracingInterceptor - - # Construct inside the try so a constructor failure (not just a missing - # contrib) also falls back to a no-op instead of aborting worker startup. - return [TracingInterceptor()] - except Exception as exc: # contrib unavailable OR constructor failure -> no-op, never raise - logger.warning( - "Temporal OTel trace interceptor unavailable (%s); traces will not propagate across Temporal boundaries.", - exc, - ) - return [] + return _sgp_temporal.interceptors() diff --git a/src/agentex/lib/core/tracing/trace.py b/src/agentex/lib/core/tracing/trace.py index 447ed7d9b..0d3f28e7a 100644 --- a/src/agentex/lib/core/tracing/trace.py +++ b/src/agentex/lib/core/tracing/trace.py @@ -12,12 +12,10 @@ from agentex.types.span import Span from agentex.lib.utils.logging import make_logger from agentex.lib.utils.model_utils import recursive_model_dump -from agentex.lib.core.tracing.obs_ids import obs_correlation from agentex.lib.core.tracing.obs_span import ( ObsSpanHandle, - open_obs_span, + begin_obs, 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 ( @@ -154,12 +152,10 @@ def _begin_obs( OTel regardless of ``SGP_OBS_MODE``, so a plain ``dd_only`` read would otherwise point at an unrelated ddtrace span). """ - if _in_tracing_dispatch_activity(): - tag_ambient_obs_span(business_span_id=span_id, business_trace_id=trace_id, prefer_otel=True) - return None, obs_correlation(prefer_otel=True) - handle = open_obs_span(name, business_span_id=span_id, business_trace_id=trace_id) - correlation = handle.correlation if handle is not None else obs_correlation() - return handle, correlation + # The wrapper-vs-ambient decision + backend selection are delegated to the + # sgp_obs Correlator (see obs_span.begin_obs); this is the single seam both the + # sync and async start_span paths share. + return begin_obs(name, span_id, trace_id) class Trace: diff --git a/src/agentex/lib/sdk/fastacp/base/base_acp_server.py b/src/agentex/lib/sdk/fastacp/base/base_acp_server.py index 1ea8e82e6..43ab0ecc1 100644 --- a/src/agentex/lib/sdk/fastacp/base/base_acp_server.py +++ b/src/agentex/lib/sdk/fastacp/base/base_acp_server.py @@ -13,6 +13,7 @@ from pydantic import TypeAdapter, ValidationError from starlette.types import Send, Scope, ASGIApp, Receive from fastapi.responses import StreamingResponse +from sgp_obs.traces.ingress import TraceContextASGIMiddleware from agentex.protocol.acp import ( RPC_SYNC_METHODS, @@ -82,6 +83,12 @@ def __init__(self): # Method handlers # this just adds a request ID to the request and response headers self.add_middleware(RequestIDMiddleware) + # Continue an inbound W3C traceparent as the active OTel context for the + # whole request (and its background Temporal dispatch), so a turn's obs + # trace stays connected instead of detaching into a fresh root. Added last + # = outermost, so it wraps RequestIDMiddleware + the handler. Delegated to + # sgp_obs; fail-open. + self.add_middleware(TraceContextASGIMiddleware) self._handlers: dict[RPCMethod, Callable] = {} # Agent info to return in healthz @@ -105,6 +112,22 @@ def _setup_handlers(self): def get_lifespan_function(self): @asynccontextmanager async def lifespan_context(app: FastAPI): # noqa: ARG001 + # Install the obs tracing provider (delegated to sgp_obs) so the per-step + # wrapper spans actually record AND export to the OTLP collector. Without + # this, get_tracer() resolves to the API-default ProxyTracerProvider — the + # forward edge (obs ids in span metadata) still fills from ambient/propagated + # context, but no obs span carrying the `.business_trace_id` reverse + # anchor is ever exported, so the business->obs pivot has nothing to land on. + # Fail-open: init_tracing never raises, but guard the import defensively too. + try: + import os + + from sgp_obs.traces import init_tracing, TracingConfig + + init_tracing(TracingConfig.from_env(service=os.getenv("OTEL_SERVICE_NAME") or "agentex-agent")) + except Exception: + logger.warning("sgp_obs tracing init skipped; obs spans will not export", exc_info=True) + env_vars = EnvironmentVariables.refresh() if env_vars.AGENTEX_BASE_URL: # Runtime SDK<->backend contract guard: fail fast if the backend is older diff --git a/uv.lock b/uv.lock index f925c2dce..7bb886dfb 100644 --- a/uv.lock +++ b/uv.lock @@ -15,7 +15,7 @@ members = [ [[package]] name = "agentex-client" -version = "0.21.0" +version = "0.22.2" source = { editable = "." } dependencies = [ { name = "anyio" }, @@ -91,7 +91,7 @@ dev = [ [[package]] name = "agentex-sdk" -version = "0.21.0" +version = "0.22.2" source = { editable = "adk" } dependencies = [ { name = "agentex-client" }, @@ -120,6 +120,7 @@ dependencies = [ { name = "rich" }, { name = "scale-gp" }, { name = "scale-gp-beta" }, + { name = "sgp-obs" }, { name = "starlette" }, { name = "temporalio" }, { name = "typer" }, @@ -156,6 +157,7 @@ requires-dist = [ { name = "rich", specifier = ">=13.9.2,<14" }, { name = "scale-gp", specifier = ">=0.1.0a59" }, { name = "scale-gp-beta", specifier = ">=0.2.0" }, + { name = "sgp-obs", specifier = "==0.2.0rc1", index = "https://scale-307185671274.d.codeartifact.us-west-2.amazonaws.com/pypi/scale-pypi/simple/" }, { name = "starlette", specifier = ">=0.49.1" }, { name = "temporalio", specifier = ">=1.26.0,<2" }, { name = "typer", specifier = ">=0.16,<0.17" }, @@ -2806,6 +2808,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2c/e0/84d284fd5268c4dcaa54dd7209d43dbe41d2120834ee7a8b245184fe13d5/scale_gp_beta-0.2.0-py3-none-any.whl", hash = "sha256:87946b4618c464711bb7c8b132540112a1a558a57b31999f93cccb5da8339643", size = 410408, upload-time = "2026-05-04T16:35:52.286Z" }, ] +[[package]] +name = "sgp-obs" +version = "0.2.0rc1" +source = { registry = "https://scale-307185671274.d.codeartifact.us-west-2.amazonaws.com/pypi/scale-pypi/simple/" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "opentelemetry-sdk" }, +] +sdist = { url = "https://scale-307185671274.d.codeartifact.us-west-2.amazonaws.com/pypi/scale-pypi/simple/sgp-obs/0.2.0rc1/sgp_obs-0.2.0rc1.tar.gz", hash = "sha512:3f63fe0f1ea4ff4a7e763c33337856f60bda512f908662922ef7d319718c02b74b1a3ca650ae3a6b705783a2ce381d37e3ebcaffbca3af0e6cc719383addee8a", size = 27057, upload-time = "2026-08-16T04:32:33.653Z" } +wheels = [ + { url = "https://scale-307185671274.d.codeartifact.us-west-2.amazonaws.com/pypi/scale-pypi/simple/sgp-obs/0.2.0rc1/sgp_obs-0.2.0rc1-py3-none-any.whl", hash = "sha512:c503912f5b72ae0ea6828f15185d8c3d01c9463fed191139db6e3b929860ca54bae68e1519854f43ec8f92c8327b7714a16bf7484f01cea2f5ae158320d9d94f", size = 39094, upload-time = "2026-08-16T04:32:33.26Z" }, +] + [[package]] name = "shellingham" version = "1.5.4" From ed04a9e63e1aee4169c48c0dc3a2c4740ce5bd6b Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Mon, 17 Aug 2026 00:26:20 -0700 Subject: [PATCH 06/12] refactor(tracing): drop backward-compat shims, use sgp-obs types directly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The obs_ids/obs_span modules were pure backward-compat adapters: they preserved the SDK's legacy surface (get_obs_mode/obs_correlation/open_obs_span/ close_obs_span/tag_ambient_obs_span/ObsSpanHandle + DD_ONLY/LGTM strings) by translating to/from sgp_obs's own types on every call. That adapter layer is the bulk of the delegation and buys nothing — trace.py is the only consumer. Delete both modules and have trace.py talk to sgp_obs directly: build the SpanRequest/BusinessRef, run the Correlator, map SGP_OBS_MODE -> ObsMode, and close via the sgp_obs handle (SpanError). Also removes the now-dead _in_tracing_dispatch_activity (the wrapper-vs-ambient decision lives entirely in the Correlator via is_dispatch_boundary). Tests: delete the shim-targeted suites (test_obs_ids/test_obs_span/ test_obs_span_fallback/test_temporal_obs_backend) — that behavior is owned by sgp_obs's own tests now. Keep the registry/error/interceptor tests (they cover trace.py's own logic); the registry test uses a minimal fake handle. Co-Authored-By: Claude Opus 4.8 --- src/agentex/lib/core/tracing/obs_ids.py | 72 ---- src/agentex/lib/core/tracing/obs_span.py | 134 ------ src/agentex/lib/core/tracing/trace.py | 103 ++--- tests/lib/core/tracing/test_obs_ids.py | 126 ------ tests/lib/core/tracing/test_obs_span.py | 516 ----------------------- tests/test_obs_handle_registry.py | 22 +- tests/test_obs_span_fallback.py | 116 ----- tests/test_temporal_obs_backend.py | 167 -------- 8 files changed, 72 insertions(+), 1184 deletions(-) delete mode 100644 src/agentex/lib/core/tracing/obs_ids.py delete mode 100644 src/agentex/lib/core/tracing/obs_span.py delete mode 100644 tests/lib/core/tracing/test_obs_ids.py delete mode 100644 tests/lib/core/tracing/test_obs_span.py delete mode 100644 tests/test_obs_span_fallback.py delete mode 100644 tests/test_temporal_obs_backend.py diff --git a/src/agentex/lib/core/tracing/obs_ids.py b/src/agentex/lib/core/tracing/obs_ids.py deleted file mode 100644 index 947bd36ba..000000000 --- a/src/agentex/lib/core/tracing/obs_ids.py +++ /dev/null @@ -1,72 +0,0 @@ -"""Correlate adk business spans with the active observability trace. - -DELEGATES to ``sgp_obs`` (the shared SGP observability library). The id-reading -that used to live here (``_lgtm_ids`` / ``_ddtrace_ids``) now lives in -``sgp_obs.traces.backends``; this module keeps the SDK's public surface -(``get_obs_mode`` / ``obs_correlation`` + the ``DD_ONLY`` / ``LGTM`` constants) so -callers are unchanged. - -Each business span is *tagged* with the active observability trace_id/span_id -(the OpenTelemetry "span link" pattern) so you can pivot a persisted business -span to its per-turn Tempo/Datadog trace, while the business trace still groups -the whole run by task id. Never fabricates ids — no active context returns ``{}``. - -Source follows ``SGP_OBS_MODE``: unset/``dd_only`` -> ddtrace; ``lgtm`` -> OTel. -Note the SDK spells the OTel mode ``"lgtm"`` while ``sgp_obs.ObsMode`` spells it -``"otel"`` — :func:`obs_mode` maps between them. -""" - -from __future__ import annotations - -import os -from typing import Dict - -from sgp_obs.traces import ObsMode -from sgp_obs.traces.backends.otel import OTelBackend -from sgp_obs.traces.backends.ddtrace import DDTraceBackend - -__all__ = ("get_obs_mode", "obs_correlation", "DD_ONLY", "LGTM") - -DD_ONLY = "dd_only" -LGTM = "lgtm" -_DEFAULT_MODE = DD_ONLY -_VALID_MODES = (DD_ONLY, LGTM) - -# Stateless singletons — the backends hold no per-call state. -_OTEL = OTelBackend() -_DDTRACE = DDTraceBackend() - - -def get_obs_mode() -> str: - """Unset/empty/unrecognized -> ``dd_only``. Returns the SDK's string form - (``"lgtm"`` / ``"dd_only"``).""" - 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 obs_mode() -> ObsMode: - """The SDK's ``SGP_OBS_MODE`` mapped to ``sgp_obs.ObsMode`` (``lgtm`` -> OTEL, - else DD_ONLY). Used to drive the sgp_obs Correlator/backends.""" - return ObsMode.OTEL if get_obs_mode() == LGTM else ObsMode.DD_ONLY - - -def obs_correlation(prefer_otel: bool = False) -> Dict[str, str]: - """Return ``{"obs_trace_id": ..., "obs_span_id": ...}`` for the active obs - context, or ``{}`` if none. Delegates id-reading to the sgp_obs backends' - ``current_ids()``. - - ``prefer_otel``: read OTel first (on the Temporal path the active span is the - temporalio OTel interceptor span regardless of ``SGP_OBS_MODE``). Never - fabricates ids — this is a correlation tag, not the span's id. - """ - try: - if prefer_otel: - corr = _OTEL.current_ids() or _DDTRACE.current_ids() - else: - corr = (_OTEL if get_obs_mode() == LGTM else _DDTRACE).current_ids() - except Exception: # obs must never fail an app call - return {} - return corr.as_metadata() if corr is not None else {} diff --git a/src/agentex/lib/core/tracing/obs_span.py b/src/agentex/lib/core/tracing/obs_span.py deleted file mode 100644 index 55742b677..000000000 --- a/src/agentex/lib/core/tracing/obs_span.py +++ /dev/null @@ -1,134 +0,0 @@ -"""Dedicated per-business-span observability wrapper span — DELEGATED to sgp_obs. - -The span/id mechanics that used to live here (``_open_otel_span`` / -``_open_ddtrace_span`` / ``_tag_*_ambient`` / ``_hex_ids``) now live in -``sgp_obs.traces.backends`` behind the ``ObsBackend`` port, and the -wrapper-vs-ambient DECISION lives in ``sgp_obs.traces.Correlator``. This module -keeps the SDK's public surface (``ObsSpanHandle`` / ``open_obs_span`` / -``close_obs_span`` / ``tag_ambient_obs_span``) as thin adapters so callers and -tests are unchanged, and adds :func:`begin_obs` which ``trace.py`` uses to -delegate the whole decision in one call. - -Reverse-tag source is ``agentex`` (this is the agentex SDK). Never raises. -""" - -from __future__ import annotations - -from typing import Dict, Optional - -from sgp_obs.traces import SpanError, Correlator, BusinessRef, SpanRequest, BusinessSource, temporal as _sgp_temporal -from sgp_obs.traces.ports import ObsSpanHandle as _SgpHandle - -from agentex.lib.core.tracing.obs_ids import LGTM, _OTEL, _DDTRACE, obs_mode, get_obs_mode - -__all__ = ( - "ObsSpanHandle", - "open_obs_span", - "close_obs_span", - "tag_ambient_obs_span", - "begin_obs", -) - - -def _to_span_error(error: Optional[Dict[str, str]]) -> Optional[SpanError]: - """SDK error dict ``{"type","message"}`` -> sgp_obs SpanError.""" - if not error: - return None - return SpanError(type=error.get("type"), message=error.get("message")) - - -def _biz(span_id: Optional[str], trace_id: Optional[str]) -> BusinessRef: - # source=agentex: this SDK is the agentex business source; the reverse tag is - # therefore ``agentex.business_span_id`` / ``agentex.business_trace_id``. - return BusinessRef(trace_id=trace_id, span_id=span_id or "", source=BusinessSource.AGENTEX) - - -def _backend(): - """The backend for the current SGP_OBS_MODE (lgtm -> OTel, else ddtrace).""" - return _OTEL if get_obs_mode() == LGTM else _DDTRACE - - -class ObsSpanHandle: - """Live handle for an open wrapper span. Adapts a sgp_obs ``ObsSpanHandle`` to - the SDK's shape: ``.correlation`` is the ``{obs_trace_id, obs_span_id}`` dict, - and ``close(error)`` takes the SDK's ``{"type","message"}`` error dict.""" - - __slots__ = ("correlation", "_inner") - - def __init__(self, inner: _SgpHandle) -> None: - self._inner = inner - self.correlation: Dict[str, str] = inner.correlation.as_metadata() - - def close(self, error: Optional[Dict[str, str]] = None) -> None: - self._inner.close(_to_span_error(error)) - - -def open_obs_span( - name: str, - business_span_id: Optional[str] = None, - business_trace_id: Optional[str] = None, -) -> Optional[ObsSpanHandle]: - """Open a step-named wrapper span in the active backend and make it active; - ``None`` when there's no live context (caller falls back to ambient ids). - Delegates to the sgp_obs backend. Never raises.""" - try: - inner = _backend().open_span(name, _biz(business_span_id, business_trace_id)) - except Exception: # pragma: no cover - backstop; obs must never break a call - return None - return ObsSpanHandle(inner) if inner is not None else None - - -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 span without opening one. - ``prefer_otel`` tags OTel first (the Temporal interceptor span is OTel - regardless of SGP_OBS_MODE). Best-effort; never raises.""" - biz = _biz(business_span_id, business_trace_id) - try: - if prefer_otel: - if _OTEL.tag_ambient(biz): - return - _DDTRACE.tag_ambient(biz) - return - _backend().tag_ambient(biz) - except Exception: # pragma: no cover - best-effort - pass - - -def close_obs_span(handle: Optional[ObsSpanHandle], error: Optional[Dict[str, str]] = None) -> None: - """Close the wrapper span, marking it errored when ``error`` is given. Safe on None.""" - if handle is None: - return - try: - handle.close(error) - except Exception: # pragma: no cover - best-effort - pass - - -def begin_obs( - name: str, - span_id: str, - trace_id: Optional[str], -) -> tuple[Optional[ObsSpanHandle], Dict[str, str]]: - """Delegate the ENTIRE wrapper-vs-ambient + backend decision to the sgp_obs - ``Correlator`` and return the SDK-shaped ``(handle | None, {obs_trace_id, - obs_span_id})``. Inside the dispatched start-span/end-span pair the Correlator - tags the ambient span (no wrapper); elsewhere it opens a per-step wrapper. - - A fresh Correlator per call keeps the SDK's per-call ``SGP_OBS_MODE`` reading - (the backends are stateless singletons; only the one-shot drift log resets). - """ - req = SpanRequest( - name=name, - business=_biz(span_id, trace_id), - in_activity=_sgp_temporal.in_activity(), - is_dispatch_boundary=_sgp_temporal.in_dispatch_boundary(), - ) - try: - inner, edge = Correlator(_OTEL, _DDTRACE, obs_mode()).begin(req) - except Exception: # pragma: no cover - obs must never break the business span - return None, {} - return (ObsSpanHandle(inner) if inner is not None else None), edge.as_metadata() diff --git a/src/agentex/lib/core/tracing/trace.py b/src/agentex/lib/core/tracing/trace.py index 0d3f28e7a..d4a859814 100644 --- a/src/agentex/lib/core/tracing/trace.py +++ b/src/agentex/lib/core/tracing/trace.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os import uuid from typing import Any, AsyncGenerator from datetime import UTC, datetime @@ -12,11 +13,18 @@ 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_span import ( - ObsSpanHandle, - begin_obs, - close_obs_span, +from sgp_obs.traces import ( + ObsMode, + SpanError, + Correlator, + BusinessRef, + SpanRequest, + BusinessSource, + temporal as _sgp_temporal, ) +from sgp_obs.traces.ports import ObsSpanHandle +from sgp_obs.traces.backends.otel import OTelBackend +from sgp_obs.traces.backends.ddtrace import DDTraceBackend from agentex.lib.core.tracing.span_error import get_span_error, set_span_error from agentex.lib.core.tracing.span_queue import ( SpanEventType, @@ -30,12 +38,34 @@ logger = make_logger(__name__) +# Stateless sgp_obs backend singletons + the SGP_OBS_MODE -> sgp_obs.ObsMode map. +# The SDK spells the OTel mode "lgtm"; sgp_obs spells it ObsMode.OTEL. +_OTEL = OTelBackend() +_DDTRACE = DDTraceBackend() + + +def _obs_mode() -> ObsMode: + """``SGP_OBS_MODE`` mapped to ``sgp_obs.ObsMode`` (``lgtm`` -> OTEL, else DD_ONLY).""" + return ObsMode.OTEL if (os.getenv("SGP_OBS_MODE") or "").strip().lower() == "lgtm" else ObsMode.DD_ONLY + + +def _close_obs(handle: ObsSpanHandle | None, error: dict[str, str] | None = None) -> None: + """Close a wrapper span (best-effort), marking it errored when the business span + carried an error. Safe on ``None``; obs must never fail the app path.""" + if handle is None: + return + try: + handle.close(SpanError(type=error.get("type"), message=error.get("message")) if error else None) + except Exception: # pragma: no cover - best-effort + pass + + # 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 +# instance can't find it, _close_obs(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. @@ -54,7 +84,7 @@ 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 + first. _close_obs 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 @@ -67,7 +97,7 @@ def _register_obs_handle(span_id: str, handle: ObsSpanHandle) -> None: _OBS_HANDLES_MAX, _evicted_id, ) - close_obs_span(evicted) + _close_obs(evicted) def _run_on_span_start(processor: SyncTracingProcessor, span: Span) -> None: @@ -104,29 +134,6 @@ def _run_on_span_end(processor: SyncTracingProcessor, span: Span) -> None: ) -def _in_tracing_dispatch_activity() -> bool: - """True only when running inside the SDK's OWN dispatched START_SPAN / END_SPAN - activity (the ``in_temporal_workflow()`` path, where a workflow runs span start - and end as SEPARATE activities that Temporal can route to different workers). - - That is the one case a per-step obs wrapper can't work: the wrapper opened in - the START_SPAN activity could never be closed by the END_SPAN activity. A span - created directly inside a *business* activity (an agent turn's own - ``adk.tracing.span``) runs start AND end in the same activity process, so a - wrapper there is safe -- it nests under the interceptor's ambient RunActivity - span and closes in-process. The tracing dispatch activities are named - ``start-span`` / ``end-span`` (``TracingActivityName``). Never raises; False - when temporalio isn't importable or we're not in an activity.""" - try: - from temporalio import activity - - if not activity.in_activity(): - return False - return activity.info().activity_type in ("start-span", "end-span") - except Exception: - return False - - def _begin_obs( name: str, span_id: str, @@ -139,23 +146,23 @@ def _begin_obs( stable/meaningful (not an arbitrary innermost httpx span), and it carries the reverse tag (business span/trace id) for the obs -> business pivot. - We open a real per-step wrapper on the sync path AND inside a *business* - Temporal activity -- there the wrapper nests under the interceptor's ambient - RunActivity span and start/end run in-process, so it closes cleanly and each - business step gets its own obs span (1:1), just like sync. - - The ONE exception is the SDK's own dispatched START_SPAN / END_SPAN activity - (a workflow calling ``adk.tracing`` -- see ``_in_tracing_dispatch_activity``): - there start and end are separate activities on possibly different workers, so - a wrapper could never be closed. We fall back to tagging the ambient - interceptor span instead, with ``prefer_otel=True`` (the interceptor span is - OTel regardless of ``SGP_OBS_MODE``, so a plain ``dd_only`` read would - otherwise point at an unrelated ddtrace span). + The whole wrapper-vs-ambient decision + backend selection is the sgp_obs + ``Correlator``'s job: on the SDK's own dispatched start-span/end-span activity + (``is_dispatch_boundary``) it tags the ambient interceptor span instead of + opening a wrapper that could never be closed; elsewhere it opens a per-step + wrapper. Fail-open — obs must never break the business span. """ - # The wrapper-vs-ambient decision + backend selection are delegated to the - # sgp_obs Correlator (see obs_span.begin_obs); this is the single seam both the - # sync and async start_span paths share. - return begin_obs(name, span_id, trace_id) + req = SpanRequest( + name=name, + business=BusinessRef(trace_id=trace_id, span_id=span_id or "", source=BusinessSource.AGENTEX), + in_activity=_sgp_temporal.in_activity(), + is_dispatch_boundary=_sgp_temporal.in_dispatch_boundary(), + ) + try: + handle, edge = Correlator(_OTEL, _DDTRACE, _obs_mode()).begin(req) + except Exception: # pragma: no cover - obs must never break the business span + return None, {} + return handle, edge.as_metadata() class Trace: @@ -258,7 +265,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(_OBS_HANDLES.pop(span.id, None), error=get_span_error(span)) + _close_obs(_OBS_HANDLES.pop(span.id, None), 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 @@ -432,7 +439,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(_OBS_HANDLES.pop(span.id, None), error=get_span_error(span)) + _close_obs(_OBS_HANDLES.pop(span.id, None), 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_ids.py b/tests/lib/core/tracing/test_obs_ids.py deleted file mode 100644 index 5cdeb81b8..000000000 --- a/tests/lib/core/tracing/test_obs_ids.py +++ /dev/null @@ -1,126 +0,0 @@ -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 deleted file mode 100644 index a7f40a511..000000000 --- a/tests/lib/core/tracing/test_obs_span.py +++ /dev/null @@ -1,516 +0,0 @@ -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_obs_handle_registry.py b/tests/test_obs_handle_registry.py index 02d3adf0a..c4ed16547 100644 --- a/tests/test_obs_handle_registry.py +++ b/tests/test_obs_handle_registry.py @@ -27,7 +27,19 @@ 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 + + +class _FakeHandle: + """Minimal stand-in for a sgp_obs ``ObsSpanHandle`` — the registry only ever + calls ``close()`` on eviction, so that's all we implement.""" + + def __init__(self, on_close: Any = None) -> None: + self.correlation: dict[str, str] = {} + self._on_close = on_close + + def close(self, error: Any = None) -> None: + if self._on_close is not None: + self._on_close(error) @pytest.fixture(autouse=True) @@ -96,8 +108,8 @@ def test_start_span_survives_raising_processor_and_no_leak(monkeypatch: pytest.M 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)) + def _make_handle(marker: str) -> _FakeHandle: + return _FakeHandle(on_close=lambda _err=None, _m=marker: closed.append(_m)) # Fill exactly to the cap: nothing evicted yet. for i in range(_OBS_HANDLES_MAX): @@ -114,8 +126,8 @@ def _make_handle(marker: str) -> ObsSpanHandle: def test_reinserting_same_span_id_refreshes_recency() -> None: - def _noop_handle() -> ObsSpanHandle: - return ObsSpanHandle(correlation={}, close=lambda _err=None: None) + def _noop_handle() -> _FakeHandle: + return _FakeHandle() trace_mod._register_obs_handle("a", _noop_handle()) trace_mod._register_obs_handle("b", _noop_handle()) diff --git a/tests/test_obs_span_fallback.py b/tests/test_obs_span_fallback.py deleted file mode 100644 index c92a42e34..000000000 --- a/tests/test_obs_span_fallback.py +++ /dev/null @@ -1,116 +0,0 @@ -"""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 deleted file mode 100644 index 1ca1f5959..000000000 --- a/tests/test_temporal_obs_backend.py +++ /dev/null @@ -1,167 +0,0 @@ -"""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. - # Option A (tag the ambient interceptor span, no wrapper) now applies only - # inside the SDK's dispatched START_SPAN/END_SPAN activity, not any activity. - monkeypatch.setenv("SGP_OBS_MODE", "dd_only") - monkeypatch.setattr(trace_mod, "_in_tracing_dispatch_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 == {} - - -class _FakeHandle: - def __init__(self, corr): - self.correlation = corr - - -def test_begin_obs_opens_wrapper_outside_dispatch_activity(monkeypatch: pytest.MonkeyPatch) -> None: - """Sync path or inside a business Temporal activity: open a per-step wrapper - (1:1), NOT Option A. Each business span gets its own obs span.""" - monkeypatch.setattr(trace_mod, "_in_tracing_dispatch_activity", lambda: False) - monkeypatch.setattr( - trace_mod, "open_obs_span", - lambda *a, **k: _FakeHandle({"obs_trace_id": "t1", "obs_span_id": "s1"}), - ) - handle, corr = trace_mod._begin_obs("mortgage.classify_intent", "bs", "bt") - assert handle is not None - assert corr == {"obs_trace_id": "t1", "obs_span_id": "s1"} - - -def test_begin_obs_tags_ambient_inside_dispatch_activity(monkeypatch: pytest.MonkeyPatch) -> None: - """Inside the dispatched START_SPAN/END_SPAN activity: no wrapper (would leak - across activities); tag the ambient interceptor span instead (Option A).""" - monkeypatch.setattr(trace_mod, "_in_tracing_dispatch_activity", lambda: True) - tagged: dict = {} - monkeypatch.setattr(trace_mod, "tag_ambient_obs_span", lambda **k: tagged.update(k)) - monkeypatch.setattr(trace_mod, "obs_correlation", lambda **k: {"obs_trace_id": "amb", "obs_span_id": "amb"}) - handle, corr = trace_mod._begin_obs("mortgage.advisor.turn", "bs", "bt") - assert handle is None - assert tagged.get("business_span_id") == "bs" and tagged.get("prefer_otel") is True - assert corr == {"obs_trace_id": "amb", "obs_span_id": "amb"} From e01fe7adbe845c6d59b05d39d58c13eebae70ba1 Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Mon, 17 Aug 2026 00:40:41 -0700 Subject: [PATCH 07/12] refactor(tracing): move the ACP->Temporal dispatch span into sgp-obs _acp_dispatch_span was per-service boundary glue living in the SDK. The ACP->Temporal dispatch is a transport boundary, not service logic, so it now lives once in the library (sgp_obs.traces.dispatch_span). Replace the 45-line function + its sys/contextlib imports with two `with dispatch_span(...)` calls, passing the business id as an attribute (agentex.task_id) so the generic library stays source-agnostic. Bumps the sgp-obs pin to 0.3.0 (adds dispatch_span + the ALWAYS_ON sampler fix). Co-Authored-By: Claude Opus 4.8 --- adk/pyproject.toml | 2 +- .../services/temporal_task_service.py | 57 ++----------------- 2 files changed, 5 insertions(+), 54 deletions(-) diff --git a/adk/pyproject.toml b/adk/pyproject.toml index 10e8a74e9..11c4fb768 100644 --- a/adk/pyproject.toml +++ b/adk/pyproject.toml @@ -62,7 +62,7 @@ dependencies = [ # SGP unified observability library: the obs edge (correlation, wrapper span, # Temporal/ingress propagation) is delegated to this shared lib. rc pin while # it stabilizes; flip to ==0.3.0 once that release publishes. - "sgp-obs==0.2.0rc1", + "sgp-obs==0.3.0", "json_log_formatter>=1.1.1", ] 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 20eb9d56e..1d37224e8 100644 --- a/src/agentex/lib/core/temporal/services/temporal_task_service.py +++ b/src/agentex/lib/core/temporal/services/temporal_task_service.py @@ -1,10 +1,9 @@ from __future__ import annotations -import sys from typing import Any from datetime import timedelta -from contextlib import contextmanager -from collections.abc import Iterator + +from sgp_obs.traces import dispatch_span from agentex.types.task import Task from agentex.types.agent import Agent @@ -16,54 +15,6 @@ 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: """ @@ -89,7 +40,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("acp.task_create", task_id=task.id): + with dispatch_span("acp.task_create", {"agentex.task_id": task.id}): return await self._temporal_client.start_workflow( workflow=self._env_vars.WORKFLOW_NAME, arg=CreateTaskParams( @@ -111,7 +62,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("acp.event_send", task_id=task.id): + with dispatch_span("acp.event_send", {"agentex.task_id": task.id}): return await self._temporal_client.send_signal( workflow_id=task.id, signal=SignalName.RECEIVE_EVENT.value, From 3aeb69d6fdec17997fc9e13a6c334c5152e98333 Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Mon, 17 Aug 2026 12:28:59 -0700 Subject: [PATCH 08/12] feat(tracing): suppress instrumentation around span export MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Now that init_tracing auto-instruments outbound httpx, the span-export calls get traced themselves and nest back into the trace being exported — PUT /v5/spans/batch (SGP) and POST /spans (Agentex), plus the egp auth/db subtree a propagated traceparent drags in. Wrap both processors' async exports in sgp_obs.traces.suppress_instrumentation so the export makes no span and injects no traceparent. Verified in sgp-dev: export-pollution spans went 51 -> 0 while each business step still shows its real work (e.g. run_agent_streamed -> the LLM call). Co-Authored-By: Claude Opus 4.8 --- .../core/tracing/processors/agentex_tracing_processor.py | 8 ++++++-- .../lib/core/tracing/processors/sgp_tracing_processor.py | 7 +++++-- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/src/agentex/lib/core/tracing/processors/agentex_tracing_processor.py b/src/agentex/lib/core/tracing/processors/agentex_tracing_processor.py index 448d013e9..dd6b3a2ef 100644 --- a/src/agentex/lib/core/tracing/processors/agentex_tracing_processor.py +++ b/src/agentex/lib/core/tracing/processors/agentex_tracing_processor.py @@ -3,6 +3,8 @@ import weakref from typing import TYPE_CHECKING, Any, Dict, override +from sgp_obs.traces import suppress_instrumentation + from agentex import Agentex from agentex.types.span import Span from agentex.lib.types.tracing import AgentexTracingProcessorConfig @@ -188,14 +190,16 @@ async def on_span_start(self, span: Span) -> None: # _skip_span_start_enabled) so each span is persisted once, on end. if self._skip_span_start: return - await self.client.spans.create(**_create_kwargs(span)) + with suppress_instrumentation(): + await self.client.spans.create(**_create_kwargs(span)) @override async def on_span_end(self, span: Span) -> None: # End-only ingest: the start create was skipped, so persist the complete # span as a single INSERT here (a bare spans.update would 404 — no row). if self._skip_span_start: - await self.client.spans.create(**_create_kwargs(span)) + with suppress_instrumentation(): + await self.client.spans.create(**_create_kwargs(span)) return update: Dict[str, Any] = {} diff --git a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py index 6d186de5f..9b4458a6d 100644 --- a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py +++ b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py @@ -9,6 +9,7 @@ from scale_gp_beta import SGPClient, AsyncSGPClient from scale_gp_beta.lib.tracing import create_span, flush_queue from scale_gp_beta.lib.tracing.span import Span as SGPSpan +from sgp_obs.traces import suppress_instrumentation from agentex.types.span import Span from agentex.lib.types.tracing import SGPTracingProcessorConfig @@ -204,7 +205,8 @@ async def on_spans_start(self, spans: list[Span]) -> None: return sgp_spans = [_build_sgp_span(span, self.env_vars) for span in spans] - await client.spans.upsert_batch(items=[s.to_request_params() for s in sgp_spans]) + with suppress_instrumentation(): + await client.spans.upsert_batch(items=[s.to_request_params() for s in sgp_spans]) _metrics.record_export_success( event_type="start", span_count=len(spans), processor="sgp" ) @@ -223,7 +225,8 @@ async def on_spans_end(self, spans: list[Span]) -> None: sgp_span = _build_sgp_span(span, self.env_vars) sgp_span.end_time = span.end_time.isoformat() # type: ignore[union-attr] sgp_spans.append(sgp_span) - await client.spans.upsert_batch(items=[s.to_request_params() for s in sgp_spans]) + with suppress_instrumentation(): + await client.spans.upsert_batch(items=[s.to_request_params() for s in sgp_spans]) _metrics.record_export_success( event_type="end", span_count=len(spans), processor="sgp" ) From 0c62db1f0f35c8818c74a083856c351266c31d8a Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Mon, 17 Aug 2026 12:29:29 -0700 Subject: [PATCH 09/12] build(tracing): pin sgp-obs[http] so httpx egress instrumentation is on by default MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every agentex agent makes outbound HTTP (LLM/egp), so httpx instrumentation is universal — pull the extra in the SDK so init_tracing's egress instrumentation (and the export suppression that pairs with it) actually functions without each agent opting in. DB instrumentation (sgp-obs[db]) stays per-agent opt-in. Co-Authored-By: Claude Opus 4.8 --- adk/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/adk/pyproject.toml b/adk/pyproject.toml index 11c4fb768..608926bfd 100644 --- a/adk/pyproject.toml +++ b/adk/pyproject.toml @@ -62,7 +62,7 @@ dependencies = [ # SGP unified observability library: the obs edge (correlation, wrapper span, # Temporal/ingress propagation) is delegated to this shared lib. rc pin while # it stabilizes; flip to ==0.3.0 once that release publishes. - "sgp-obs==0.3.0", + "sgp-obs[http]==0.3.0", "json_log_formatter>=1.1.1", ] From 5836f5ece5c2e15bee971f8c6e11e786ba23d1c5 Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Mon, 17 Aug 2026 16:44:05 -0700 Subject: [PATCH 10/12] feat(llm): instrument litellm streaming via sgp-obs instrument_stream MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wrap both completion_stream (sync) and acompletion_stream (async) with sgp_obs.traces.instrument_stream so every litellm streaming call emits one gen_ai.chat span with TTFT/TTAT events + the decode-window timing (tps, tpot, output tokens) — filling the "span for TTFT but nothing for the streaming" gap for non-Temporal / litellm agents. Provider extractors count content-bearing deltas; fail-open, chunks pass through unchanged. Co-Authored-By: Claude Opus 4.8 --- .../lib/core/adapters/llm/adapter_litellm.py | 65 +++++++++++++++++-- 1 file changed, 59 insertions(+), 6 deletions(-) diff --git a/src/agentex/lib/core/adapters/llm/adapter_litellm.py b/src/agentex/lib/core/adapters/llm/adapter_litellm.py index 7935f5f49..3235e14f1 100644 --- a/src/agentex/lib/core/adapters/llm/adapter_litellm.py +++ b/src/agentex/lib/core/adapters/llm/adapter_litellm.py @@ -1,7 +1,8 @@ -from typing import override -from collections.abc import Generator, AsyncGenerator +from typing import Any, override +from collections.abc import Mapping, Generator, AsyncGenerator import litellm as llm +from sgp_obs.traces import instrument_stream, instrument_stream_sync from agentex.lib.utils.logging import make_logger from agentex.lib.types.llm_messages import Completion @@ -10,6 +11,37 @@ logger = make_logger(__name__) +def _delta_content(chunk: Completion) -> Any: + """Text delta of a streaming chunk, or ``None`` — defensive against provider shape. + Obs extractors must never raise into the stream.""" + try: + choices = getattr(chunk, "choices", None) or [] + if not choices: + return None + delta = getattr(choices[0], "delta", None) + return getattr(delta, "content", None) if delta is not None else None + except Exception: + return None + + +def _output_tokens(chunk: Completion) -> int: + # 1 per content-bearing delta — a good streaming proxy without a tokenizer. + return 1 if _delta_content(chunk) else 0 + + +def _is_answer(chunk: Completion) -> bool: + # First user-visible answer token (text); skips role-only / tool-call / empty deltas. + return bool(_delta_content(chunk)) + + +def _stream_attrs(args: tuple[Any, ...], kwargs: dict[str, Any]) -> Mapping[str, Any]: + model = kwargs.get("model") or (args[0] if args else None) + attrs: dict[str, Any] = {"gen_ai.system": "litellm", "gen_ai.operation.name": "chat"} + if model: + attrs["gen_ai.request.model"] = str(model) + return attrs + + class LiteLLMGateway(LLMGateway): @override def completion(self, *args, **kwargs) -> Completion: @@ -26,8 +58,19 @@ def completion_stream(self, *args, **kwargs) -> Generator[Completion, None, None if not kwargs.get("stream"): raise ValueError("To use streaming, please set stream=True in the kwargs") - for chunk in llm.completion(*args, **kwargs): - yield Completion.model_validate(chunk) + def _chunks() -> Generator[Completion, None, None]: + for chunk in llm.completion(*args, **kwargs): + yield Completion.model_validate(chunk) + + # Wrap the whole generation in one gen_ai.chat span (TTFT/TTAT events + + # decode-window tps/tpot/output-tokens). Fail-open; chunks pass through. + yield from instrument_stream_sync( + _chunks(), + name="gen_ai.chat", + attributes=_stream_attrs(args, kwargs), + output_tokens=_output_tokens, + is_answer=_is_answer, + ) @override async def acompletion(self, *args, **kwargs) -> Completion: @@ -47,5 +90,15 @@ async def acompletion_stream( if not kwargs.get("stream"): raise ValueError("To use streaming, please set stream=True in the kwargs") - async for chunk in await llm.acompletion(*args, **kwargs): # type: ignore[misc] - yield Completion.model_validate(chunk) + async def _chunks() -> AsyncGenerator[Completion, None]: + async for chunk in await llm.acompletion(*args, **kwargs): # type: ignore[misc] + yield Completion.model_validate(chunk) + + async for completion in instrument_stream( + _chunks(), + name="gen_ai.chat", + attributes=_stream_attrs(args, kwargs), + output_tokens=_output_tokens, + is_answer=_is_answer, + ): + yield completion From b321c9d7206faaf664c5e021bb01fd3928595144 Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Mon, 17 Aug 2026 16:47:26 -0700 Subject: [PATCH 11/12] feat(llm): instrument scale-gp streaming; share stream-obs extractors Wire SGPLLMGateway.{completion_stream,acompletion_stream} through sgp_obs.traces.instrument_stream so scale-gp streaming gets the same gen_ai.chat span (TTFT/TTAT + decode-window timing) as litellm. Both gateways yield the same OpenAI-shaped Completion chunk, so the extractors + attribute builder move to a shared _stream_obs module (litellm refactored to use it; scale-gp tags gen_ai.system=scale-gp). Streaming instrumentation is now uniform across adapters. Co-Authored-By: Claude Opus 4.8 --- .../lib/core/adapters/llm/_stream_obs.py | 46 ++++++++++++++++++ .../lib/core/adapters/llm/adapter_litellm.py | 48 ++++--------------- .../lib/core/adapters/llm/adapter_sgp.py | 31 ++++++++++-- 3 files changed, 82 insertions(+), 43 deletions(-) create mode 100644 src/agentex/lib/core/adapters/llm/_stream_obs.py diff --git a/src/agentex/lib/core/adapters/llm/_stream_obs.py b/src/agentex/lib/core/adapters/llm/_stream_obs.py new file mode 100644 index 000000000..5be4e0106 --- /dev/null +++ b/src/agentex/lib/core/adapters/llm/_stream_obs.py @@ -0,0 +1,46 @@ +"""Shared GenAI-streaming instrumentation for the LLM adapters. + +Both the litellm and scale-gp gateways yield the same OpenAI-shaped ``Completion`` +chunk, so one set of extractors + attribute builders drives +``sgp_obs.traces.instrument_stream`` for both. Every extractor is defensive — an +obs helper must never raise into the token stream. +""" + +from __future__ import annotations + +from typing import Any +from collections.abc import Mapping + +from agentex.lib.types.llm_messages import Completion + + +def delta_content(chunk: Completion) -> Any: + """Text delta of a streaming chunk, or ``None`` — tolerant of provider shape.""" + try: + choices = getattr(chunk, "choices", None) or [] + if not choices: + return None + delta = getattr(choices[0], "delta", None) + return getattr(delta, "content", None) if delta is not None else None + except Exception: + return None + + +def output_tokens(chunk: Completion) -> int: + # 1 per content-bearing delta — a good streaming proxy without a tokenizer. + return 1 if delta_content(chunk) else 0 + + +def is_answer(chunk: Completion) -> bool: + # First user-visible answer token (text); skips role-only / tool-call / empty deltas. + return bool(delta_content(chunk)) + + +def stream_attrs(system: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Mapping[str, Any]: + """GenAI-semconv attributes for the ``gen_ai.chat`` span. ``model`` is read from + kwargs or the first positional arg, matching how the gateways call the client.""" + model = kwargs.get("model") or (args[0] if args else None) + attrs: dict[str, Any] = {"gen_ai.system": system, "gen_ai.operation.name": "chat"} + if model: + attrs["gen_ai.request.model"] = str(model) + return attrs diff --git a/src/agentex/lib/core/adapters/llm/adapter_litellm.py b/src/agentex/lib/core/adapters/llm/adapter_litellm.py index 3235e14f1..7e0c7b0a5 100644 --- a/src/agentex/lib/core/adapters/llm/adapter_litellm.py +++ b/src/agentex/lib/core/adapters/llm/adapter_litellm.py @@ -1,5 +1,5 @@ -from typing import Any, override -from collections.abc import Mapping, Generator, AsyncGenerator +from typing import override +from collections.abc import Generator, AsyncGenerator import litellm as llm from sgp_obs.traces import instrument_stream, instrument_stream_sync @@ -7,41 +7,11 @@ from agentex.lib.utils.logging import make_logger from agentex.lib.types.llm_messages import Completion from agentex.lib.core.adapters.llm.port import LLMGateway +from agentex.lib.core.adapters.llm._stream_obs import is_answer, stream_attrs, output_tokens logger = make_logger(__name__) -def _delta_content(chunk: Completion) -> Any: - """Text delta of a streaming chunk, or ``None`` — defensive against provider shape. - Obs extractors must never raise into the stream.""" - try: - choices = getattr(chunk, "choices", None) or [] - if not choices: - return None - delta = getattr(choices[0], "delta", None) - return getattr(delta, "content", None) if delta is not None else None - except Exception: - return None - - -def _output_tokens(chunk: Completion) -> int: - # 1 per content-bearing delta — a good streaming proxy without a tokenizer. - return 1 if _delta_content(chunk) else 0 - - -def _is_answer(chunk: Completion) -> bool: - # First user-visible answer token (text); skips role-only / tool-call / empty deltas. - return bool(_delta_content(chunk)) - - -def _stream_attrs(args: tuple[Any, ...], kwargs: dict[str, Any]) -> Mapping[str, Any]: - model = kwargs.get("model") or (args[0] if args else None) - attrs: dict[str, Any] = {"gen_ai.system": "litellm", "gen_ai.operation.name": "chat"} - if model: - attrs["gen_ai.request.model"] = str(model) - return attrs - - class LiteLLMGateway(LLMGateway): @override def completion(self, *args, **kwargs) -> Completion: @@ -67,9 +37,9 @@ def _chunks() -> Generator[Completion, None, None]: yield from instrument_stream_sync( _chunks(), name="gen_ai.chat", - attributes=_stream_attrs(args, kwargs), - output_tokens=_output_tokens, - is_answer=_is_answer, + attributes=stream_attrs("litellm", args, kwargs), + output_tokens=output_tokens, + is_answer=is_answer, ) @override @@ -97,8 +67,8 @@ async def _chunks() -> AsyncGenerator[Completion, None]: async for completion in instrument_stream( _chunks(), name="gen_ai.chat", - attributes=_stream_attrs(args, kwargs), - output_tokens=_output_tokens, - is_answer=_is_answer, + attributes=stream_attrs("litellm", args, kwargs), + output_tokens=output_tokens, + is_answer=is_answer, ): yield completion diff --git a/src/agentex/lib/core/adapters/llm/adapter_sgp.py b/src/agentex/lib/core/adapters/llm/adapter_sgp.py index 31098246e..423eae70d 100644 --- a/src/agentex/lib/core/adapters/llm/adapter_sgp.py +++ b/src/agentex/lib/core/adapters/llm/adapter_sgp.py @@ -5,10 +5,12 @@ from collections.abc import Generator, AsyncGenerator from scale_gp import SGPClient, AsyncSGPClient +from sgp_obs.traces import instrument_stream, instrument_stream_sync from agentex.lib.utils.logging import make_logger from agentex.lib.types.llm_messages import Completion from agentex.lib.core.adapters.llm.port import LLMGateway +from agentex.lib.core.adapters.llm._stream_obs import is_answer, stream_attrs, output_tokens logger = make_logger(__name__) @@ -35,8 +37,19 @@ def completion_stream(self, *args, **kwargs) -> Generator[Completion, None, None if not kwargs.get("stream"): raise ValueError("To use streaming, please set stream=True in the kwargs") - for chunk in self.sync_client.beta.chat.completions.create(*args, **kwargs): - yield Completion.model_validate(chunk) + def _chunks() -> Generator[Completion, None, None]: + for chunk in self.sync_client.beta.chat.completions.create(*args, **kwargs): + yield Completion.model_validate(chunk) + + # Wrap the whole generation in one gen_ai.chat span (TTFT/TTAT events + + # decode-window tps/tpot/output-tokens). Fail-open; chunks pass through. + yield from instrument_stream_sync( + _chunks(), + name="gen_ai.chat", + attributes=stream_attrs("scale-gp", args, kwargs), + output_tokens=output_tokens, + is_answer=is_answer, + ) @override async def acompletion(self, *args, **kwargs) -> Completion: @@ -56,5 +69,15 @@ async def acompletion_stream( if not kwargs.get("stream"): raise ValueError("To use streaming, please set stream=True in the kwargs") - async for chunk in self.async_client.beta.chat.completions.create(*args, **kwargs): # type: ignore[misc] - yield Completion.model_validate(chunk) + async def _chunks() -> AsyncGenerator[Completion, None]: + async for chunk in self.async_client.beta.chat.completions.create(*args, **kwargs): # type: ignore[misc] + yield Completion.model_validate(chunk) + + async for completion in instrument_stream( + _chunks(), + name="gen_ai.chat", + attributes=stream_attrs("scale-gp", args, kwargs), + output_tokens=output_tokens, + is_answer=is_answer, + ): + yield completion From 0f72af42adfcff1a19c88ebaece94df787b92dca Mon Sep 17 00:00:00 2001 From: Nitesh Dhanpal Date: Mon, 17 Aug 2026 21:19:25 -0700 Subject: [PATCH 12/12] feat(tracing): install the openai-agents OTel bridge at ACP startup After init_tracing, register sgp_obs's openai-agents tracing bridge so a Runner turn's internal phases (generation/tool/handoff) show up in Tempo under the business step instead of a dark gap. Fail-open no-op for non-openai-agents agents. Co-Authored-By: Claude Opus 4.8 --- src/agentex/lib/sdk/fastacp/base/base_acp_server.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/agentex/lib/sdk/fastacp/base/base_acp_server.py b/src/agentex/lib/sdk/fastacp/base/base_acp_server.py index 43ab0ecc1..a89b6ff1e 100644 --- a/src/agentex/lib/sdk/fastacp/base/base_acp_server.py +++ b/src/agentex/lib/sdk/fastacp/base/base_acp_server.py @@ -122,9 +122,13 @@ async def lifespan_context(app: FastAPI): # noqa: ARG001 try: import os - from sgp_obs.traces import init_tracing, TracingConfig + from sgp_obs.traces import init_tracing, TracingConfig, install_openai_agents_bridge init_tracing(TracingConfig.from_env(service=os.getenv("OTEL_SERVICE_NAME") or "agentex-agent")) + # Bridge openai-agents' own tracing (Runner generations/tools/handoffs) + # into OTel so a Runner turn's internals show up under the business step + # instead of a dark gap. No-op when the agents SDK isn't in use. + install_openai_agents_bridge() except Exception: logger.warning("sgp_obs tracing init skipped; obs spans will not export", exc_info=True)