diff --git a/.release-please-manifest.json b/.release-please-manifest.json index 6e2eef653..1e396bd29 100644 --- a/.release-please-manifest.json +++ b/.release-please-manifest.json @@ -1,4 +1,4 @@ { - ".": "0.22.2", - "adk": "0.22.2" + ".": "0.23.0", + "adk": "0.23.0" } 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 diff --git a/CHANGELOG.md b/CHANGELOG.md index 9b175a931..f2291964c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,21 @@ * **tracing:** emit OTel metrics for async span queue depth, batch drain, and SGP export success/failure (HTTP status labels). Disable SDK-side recording with ``AGENTEX_TRACING_METRICS=0``. +## 0.23.0 (2026-08-07) + +Full Changelog: [agentex-client-v0.22.2...agentex-client-v0.23.0](https://github.com/scaleapi/scale-agentex-python/compare/agentex-client-v0.22.2...agentex-client-v0.23.0) + +### Features + +* propagate error categories to SGP spans ([#486](https://github.com/scaleapi/scale-agentex-python/issues/486)) ([f2b1808](https://github.com/scaleapi/scale-agentex-python/commit/f2b18087c5049f7e2a159f8e1ad3bf069b01eb23)) +* **tracing:** correlate business spans with obs via dedicated wrapper span ([#484](https://github.com/scaleapi/scale-agentex-python/issues/484)) ([72732b7](https://github.com/scaleapi/scale-agentex-python/commit/72732b7c07700df840a2308424154f11a30e39f2)) +* **tracing:** propagate OTel trace context across Temporal boundaries ([#485](https://github.com/scaleapi/scale-agentex-python/issues/485)) ([da7ea15](https://github.com/scaleapi/scale-agentex-python/commit/da7ea1558683da05a3f9ecb119b91bf873437be1)) + + +### Bug Fixes + +* **task-create:** handle WorkflowAlreadyStartedError gracefully ([#489](https://github.com/scaleapi/scale-agentex-python/issues/489)) ([462195d](https://github.com/scaleapi/scale-agentex-python/commit/462195d11c48713f1b15487adad45776043a09fd)) + ## 0.22.2 (2026-07-30) Full Changelog: [agentex-client-v0.22.1...agentex-client-v0.22.2](https://github.com/scaleapi/scale-agentex-python/compare/agentex-client-v0.22.1...agentex-client-v0.22.2) diff --git a/adk/CHANGELOG.md b/adk/CHANGELOG.md index aba187304..d6451a109 100644 --- a/adk/CHANGELOG.md +++ b/adk/CHANGELOG.md @@ -1,5 +1,13 @@ # Changelog +## 0.23.0 (2026-08-07) + +Full Changelog: [agentex-sdk-v0.22.2...agentex-sdk-v0.23.0](https://github.com/scaleapi/scale-agentex-python/compare/agentex-sdk-v0.22.2...agentex-sdk-v0.23.0) + +### Features + +* propagate error categories to SGP spans ([#486](https://github.com/scaleapi/scale-agentex-python/issues/486)) ([f2b1808](https://github.com/scaleapi/scale-agentex-python/commit/f2b18087c5049f7e2a159f8e1ad3bf069b01eb23)) + ## 0.22.2 (2026-07-30) Full Changelog: [agentex-sdk-v0.22.1...agentex-sdk-v0.22.2](https://github.com/scaleapi/scale-agentex-python/compare/agentex-sdk-v0.22.1...agentex-sdk-v0.22.2) diff --git a/adk/pyproject.toml b/adk/pyproject.toml index 3569a8e52..1ec47d7c5 100644 --- a/adk/pyproject.toml +++ b/adk/pyproject.toml @@ -4,7 +4,7 @@ # (agentex/{__init__.py, _*.py, types/, resources/}) ships from the slim # sibling package `agentex-client` which is pinned as a runtime dep. name = "agentex-sdk" -version = "0.22.2" +version = "0.23.0" description = "Agent Development Kit (ADK) overlay for the Agentex API — FastACP server, Temporal workflows, LLM provider integrations, observability" license = "Apache-2.0" authors = [ @@ -53,7 +53,7 @@ dependencies = [ "pydantic-ai-slim>=1.0,<2", "langgraph-checkpoint>=2.0.0", "scale-gp>=0.1.0a59", - "scale-gp-beta>=0.2.0", + "scale-gp-beta>=0.5.0", "mcp>=1.4.1", # Observability "ddtrace>=3.13.0", diff --git a/pyproject.toml b/pyproject.toml index f73829573..d2cf0e23b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ # overlay (formerly `src/agentex/lib/*`) now lives in `adk/` and ships # as the sibling `agentex-sdk` package — see `adk/pyproject.toml`. name = "agentex-client" -version = "0.22.2" +version = "0.23.0" description = "The official Python REST client for the Agentex API" dynamic = ["readme"] license = "Apache-2.0" diff --git a/src/agentex/_version.py b/src/agentex/_version.py index 11f3825fe..25c06dc1a 100644 --- a/src/agentex/_version.py +++ b/src/agentex/_version.py @@ -1,4 +1,4 @@ # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details. __title__ = "agentex" -__version__ = "0.22.2" # x-release-please-version +__version__ = "0.23.0" # x-release-please-version diff --git a/src/agentex/lib/adk/_modules/tracing.py b/src/agentex/lib/adk/_modules/tracing.py index 7d49bb91c..9b89d076e 100644 --- a/src/agentex/lib/adk/_modules/tracing.py +++ b/src/agentex/lib/adk/_modules/tracing.py @@ -19,6 +19,7 @@ StartSpanParams, TracingActivityName, ) +from agentex.lib.core.tracing.span_error import set_span_error from agentex.lib.core.tracing.tracer import AsyncTracer from agentex.lib.core.harness.types import TurnUsage from agentex.types.span import Span @@ -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/clients/temporal/temporal_client.py b/src/agentex/lib/core/clients/temporal/temporal_client.py index 02c44d81e..8b74ddf77 100644 --- a/src/agentex/lib/core/clients/temporal/temporal_client.py +++ b/src/agentex/lib/core/clients/temporal/temporal_client.py @@ -5,7 +5,11 @@ from collections.abc import Callable from temporalio.client import Client, WorkflowExecutionStatus -from temporalio.common import RetryPolicy as TemporalRetryPolicy, WorkflowIDReusePolicy +from temporalio.common import ( + RetryPolicy as TemporalRetryPolicy, + WorkflowIDReusePolicy, + WorkflowIDConflictPolicy, +) from temporalio.service import RPCError, RPCStatusCode from temporalio.converter import PayloadCodec, DataConverter @@ -15,6 +19,7 @@ TaskStatus, RetryPolicy, WorkflowState, + ConflictWorkflowPolicy, DuplicateWorkflowPolicy, ) from agentex.lib.core.clients.temporal.utils import get_temporal_client @@ -75,6 +80,13 @@ DuplicateWorkflowPolicy.TERMINATE_IF_RUNNING: WorkflowIDReusePolicy.TERMINATE_IF_RUNNING, } +CONFLICT_POLICY_TO_ID_CONFLICT_POLICY = { + ConflictWorkflowPolicy.UNSPECIFIED: WorkflowIDConflictPolicy.UNSPECIFIED, + ConflictWorkflowPolicy.FAIL: WorkflowIDConflictPolicy.FAIL, + ConflictWorkflowPolicy.USE_EXISTING: WorkflowIDConflictPolicy.USE_EXISTING, + ConflictWorkflowPolicy.TERMINATE_EXISTING: WorkflowIDConflictPolicy.TERMINATE_EXISTING, +} + class TemporalClient: def __init__( @@ -151,11 +163,20 @@ async def start_workflow( self, *args: Any, duplicate_policy: DuplicateWorkflowPolicy = DuplicateWorkflowPolicy.ALLOW_DUPLICATE, + conflict_policy: ConflictWorkflowPolicy = ConflictWorkflowPolicy.UNSPECIFIED, retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY, task_timeout: timedelta = timedelta(seconds=10), execution_timeout: timedelta | None = None, **kwargs: Any, ) -> str: + if ( + duplicate_policy == DuplicateWorkflowPolicy.TERMINATE_IF_RUNNING + and conflict_policy != ConflictWorkflowPolicy.UNSPECIFIED + ): + raise ValueError( + "conflict_policy cannot be set when duplicate_policy is TERMINATE_IF_RUNNING; " + "use ConflictWorkflowPolicy.TERMINATE_EXISTING instead" + ) temporal_retry_policy = TemporalRetryPolicy(**retry_policy.model_dump(exclude_unset=True)) workflow_handle = await self.client.start_workflow( *args, @@ -163,6 +184,7 @@ async def start_workflow( task_timeout=task_timeout, execution_timeout=execution_timeout, id_reuse_policy=DUPLICATE_POLICY_TO_ID_REUSE_POLICY[duplicate_policy], + id_conflict_policy=CONFLICT_POLICY_TO_ID_CONFLICT_POLICY[conflict_policy], **kwargs, ) return workflow_handle.id diff --git a/src/agentex/lib/core/clients/temporal/types.py b/src/agentex/lib/core/clients/temporal/types.py index 8ce596d77..eceef154e 100644 --- a/src/agentex/lib/core/clients/temporal/types.py +++ b/src/agentex/lib/core/clients/temporal/types.py @@ -40,6 +40,13 @@ class DuplicateWorkflowPolicy(str, Enum): TERMINATE_IF_RUNNING = "TERMINATE_IF_RUNNING" +class ConflictWorkflowPolicy(str, Enum): + UNSPECIFIED = "UNSPECIFIED" + FAIL = "FAIL" + USE_EXISTING = "USE_EXISTING" + TERMINATE_EXISTING = "TERMINATE_EXISTING" + + class TaskStatus(str, Enum): CANCELED = "CANCELED" COMPLETED = "COMPLETED" 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/services/temporal_task_service.py b/src/agentex/lib/core/temporal/services/temporal_task_service.py index 5f6c0c381..774de4d9e 100644 --- a/src/agentex/lib/core/temporal/services/temporal_task_service.py +++ b/src/agentex/lib/core/temporal/services/temporal_task_service.py @@ -1,18 +1,70 @@ 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 from agentex.types.event import Event from agentex.protocol.acp import SendEventParams, CreateTaskParams, InterruptTaskParams from agentex.lib.environment_variables import EnvironmentVariables -from agentex.lib.core.clients.temporal.types import WorkflowState +from agentex.lib.core.clients.temporal.types import WorkflowState, ConflictWorkflowPolicy from agentex.lib.core.temporal.types.workflow import SignalName 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,22 @@ 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 + # USE_EXISTING makes task/create idempotent + # If same task ID is already running Temporal returns a handle to the existing run instead of raising WorkflowAlreadyStarted + 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, + conflict_policy=ConflictWorkflowPolicy.USE_EXISTING, + ) async def get_state(self, task_id: str) -> WorkflowState: """ @@ -63,16 +114,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/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/__init__.py b/src/agentex/lib/core/tracing/__init__.py index 639f3ba8e..580b53c20 100644 --- a/src/agentex/lib/core/tracing/__init__.py +++ b/src/agentex/lib/core/tracing/__init__.py @@ -1,6 +1,12 @@ from agentex.types.span import Span from agentex.lib.core.tracing.trace import Trace, AsyncTrace from agentex.lib.core.tracing.tracer import Tracer, AsyncTracer +from agentex.lib.core.tracing.span_error import ( + ErrorCategory, + PlatformError, + ApplicationError, + CategorizedError, +) from agentex.lib.core.tracing.span_queue import ( AsyncSpanQueue, get_default_span_queue, @@ -13,6 +19,10 @@ "Span", "Tracer", "AsyncTracer", + "CategorizedError", + "ApplicationError", + "PlatformError", + "ErrorCategory", "AsyncSpanQueue", "get_default_span_queue", "shutdown_default_span_queue", 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/processors/sgp_tracing_processor.py b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py index 6d186de5f..b2e4563f2 100644 --- a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py +++ b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py @@ -87,6 +87,7 @@ def _build_sgp_span(span: Span, env_vars: EnvironmentVariables) -> SGPSpan: error = get_span_error(span) if error is not None: sgp_span.set_error(error_type=error["type"], error_message=error["message"]) + sgp_span.metadata["error_category"] = error.get("category", "unknown") return sgp_span diff --git a/src/agentex/lib/core/tracing/span_error.py b/src/agentex/lib/core/tracing/span_error.py index 508c5e800..f20eae471 100644 --- a/src/agentex/lib/core/tracing/span_error.py +++ b/src/agentex/lib/core/tracing/span_error.py @@ -1,6 +1,13 @@ from __future__ import annotations -from typing import Any +from typing import Any, cast + +from scale_gp_beta.lib.tracing import ( + PlatformError as PlatformError, + ApplicationError as ApplicationError, + CategorizedError, +) +from scale_gp_beta.lib.tracing.types import ErrorCategory from agentex.types.span import Span @@ -13,14 +20,48 @@ # SGP and agentex-native span stores. SPAN_ERROR_KEY = "__error__" +ERROR_CATEGORY_UNKNOWN: ErrorCategory = "unknown" +_ERROR_CATEGORIES = frozenset({"application", "platform", "unknown"}) + + +def _normalize_error_category(value: object) -> ErrorCategory | None: + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in _ERROR_CATEGORIES: + return cast(ErrorCategory, normalized) + return None + + +def _error_category( + exc: BaseException, + explicit_category: ErrorCategory | str | None = None, +) -> ErrorCategory: + """Return an explicit producer classification, defaulting safely to unknown.""" + return ( + _normalize_error_category(explicit_category) + or (exc.error_category if isinstance(exc, CategorizedError) else None) + or ERROR_CATEGORY_UNKNOWN + ) + -def set_span_error(span: Span, exc: BaseException) -> None: +def set_span_error( + span: Span, + exc: BaseException, + *, + error_category: ErrorCategory | str | None = None, +) -> None: """Record an exception on ``span`` under ``data[SPAN_ERROR_KEY]``. + An explicit ``error_category`` takes precedence over a ``CategorizedError`` + classification. Invalid or absent categories become unknown. No-op when ``span.data`` is a list (matching ``_add_source_to_span``, which only attaches metadata to dict-shaped data). """ - error = {"type": type(exc).__name__, "message": str(exc)} + error = { + "type": type(exc).__name__, + "message": str(exc), + "category": _error_category(exc, error_category), + } if span.data is None: span.data = {} if isinstance(span.data, dict): 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/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/adk/test_tracing_module.py b/tests/lib/adk/test_tracing_module.py index 00e4aae65..c17ff5ff6 100644 --- a/tests/lib/adk/test_tracing_module.py +++ b/tests/lib/adk/test_tracing_module.py @@ -10,6 +10,7 @@ from agentex.types.span import Span from agentex.lib.core.harness.types import TurnUsage from agentex.lib.adk._modules.tracing import TurnSpan, TracingModule +from agentex.lib.core.tracing.span_error import get_span_error from agentex.lib.core.services.adk.tracing import TracingService @@ -249,6 +250,24 @@ async def test_span_context_manager_forwards_task_id(self): assert mock_service.start_span.call_args.kwargs["task_id"] == "task-abc" mock_service.end_span.assert_called_once() + async def test_span_context_manager_records_and_reraises_body_error(self): + mock_service, module = _make_module() + started = _make_span() + mock_service.start_span.return_value = started + mock_service.end_span.return_value = started + + with patch.object(_tracing_mod, "in_temporal_workflow", return_value=False): + with pytest.raises(RuntimeError, match="boom"): + async with module.span(trace_id="trace-123", name="test-span"): + raise RuntimeError("boom") + + assert get_span_error(started) == { + "type": "RuntimeError", + "message": "boom", + "category": "unknown", + } + mock_service.end_span.assert_called_once_with(trace_id="trace-123", span=started) + async def test_span_context_manager_noop_when_no_trace_id(self): mock_service, module = _make_module() diff --git a/tests/lib/core/services/test_temporal_task_service.py b/tests/lib/core/services/test_temporal_task_service.py new file mode 100644 index 000000000..7589863cd --- /dev/null +++ b/tests/lib/core/services/test_temporal_task_service.py @@ -0,0 +1,140 @@ +"""Unit tests for TemporalTaskService idempotency behavior. + +Covers the ``task/create`` idempotency guarantee: duplicate submits for the +same task ID must not raise ``WorkflowAlreadyStartedError``. The service +achieves this by passing ``ConflictWorkflowPolicy.USE_EXISTING`` through the +``TemporalClient`` wrapper, which maps to ``WorkflowIDConflictPolicy.USE_EXISTING`` +on the underlying temporalio client so Temporal returns a handle to the +existing run instead of erroring. +""" + +from __future__ import annotations + +from unittest.mock import Mock, AsyncMock + +import pytest +from temporalio.common import WorkflowIDConflictPolicy + +from agentex.types.task import Task +from agentex.types.agent import Agent +from agentex.lib.core.clients.temporal.types import ( + ConflictWorkflowPolicy, + DuplicateWorkflowPolicy, +) +from agentex.lib.core.clients.temporal.temporal_client import TemporalClient +from agentex.lib.core.temporal.services.temporal_task_service import TemporalTaskService + + +def _agent() -> Agent: + return Agent( + id="test-agent-456", + name="test-agent", + description="test-agent", + acp_type="async", + created_at="2023-01-01T00:00:00Z", + updated_at="2023-01-01T00:00:00Z", + ) + + +def _task() -> Task: + return Task(id="test-task-123", status="RUNNING") + + +def _env_vars() -> Mock: + env_vars = Mock() + env_vars.WORKFLOW_NAME = "test-workflow" + env_vars.WORKFLOW_TASK_QUEUE = "test-queue" + env_vars.WORKFLOW_EXECUTION_TIMEOUT_SECONDS = 0 + return env_vars + + +class TestSubmitTaskIdempotency: + async def test_submit_task_uses_use_existing_conflict_policy(self) -> None: + """Duplicate task/create must be idempotent. + + Passing ``ConflictWorkflowPolicy.USE_EXISTING`` tells Temporal to + return the existing workflow handle instead of raising + ``WorkflowAlreadyStartedError`` when a run with that ID is already + active. Without this, load-balanced agentex-agent replicas racing on + the same task ID surface Temporal's start conflict as an error log. + """ + temporal_client = Mock() + temporal_client.start_workflow = AsyncMock(return_value="test-task-123") + + service = TemporalTaskService(temporal_client=temporal_client, env_vars=_env_vars()) + + result = await service.submit_task(agent=_agent(), task=_task(), params=None) + + temporal_client.start_workflow.assert_awaited_once() + kwargs = temporal_client.start_workflow.await_args.kwargs + assert kwargs["conflict_policy"] == ConflictWorkflowPolicy.USE_EXISTING + assert kwargs["id"] == "test-task-123" + assert result == "test-task-123" + + +class TestTemporalClientConflictPolicyPlumbing: + """Boundary tests: ``TemporalClient.start_workflow`` wraps + ``WorkflowIDConflictPolicy`` in a local ``ConflictWorkflowPolicy`` enum + (mirroring the existing ``DuplicateWorkflowPolicy`` pattern) so SDK users + don't have to import ``temporalio.common`` to opt into non-default behavior. + Also guards the incompatible ``TERMINATE_IF_RUNNING`` + explicit-conflict + combo client-side rather than letting it round-trip to the frontend as + ``InvalidArgument``. + """ + + async def test_forwards_conflict_policy_when_set(self) -> None: + inner_client = Mock() + inner_handle = Mock() + inner_handle.id = "wf-1" + inner_client.start_workflow = AsyncMock(return_value=inner_handle) + + tc = TemporalClient(temporal_client=inner_client) + + await tc.start_workflow( + workflow="w", + arg={}, + id="id-1", + task_queue="q", + conflict_policy=ConflictWorkflowPolicy.USE_EXISTING, + ) + + kwargs = inner_client.start_workflow.await_args.kwargs + assert kwargs["id_conflict_policy"] == WorkflowIDConflictPolicy.USE_EXISTING + + async def test_default_conflict_policy_is_unspecified(self) -> None: + inner_client = Mock() + inner_handle = Mock() + inner_handle.id = "wf-1" + inner_client.start_workflow = AsyncMock(return_value=inner_handle) + + tc = TemporalClient(temporal_client=inner_client) + + await tc.start_workflow(workflow="w", arg={}, id="id-1", task_queue="q") + + kwargs = inner_client.start_workflow.await_args.kwargs + assert kwargs["id_conflict_policy"] == WorkflowIDConflictPolicy.UNSPECIFIED + + async def test_terminate_if_running_with_explicit_conflict_policy_raises(self) -> None: + """temporalio rejects this combo at the frontend as InvalidArgument; + we fail fast client-side with a clearer message. + """ + inner_client = Mock() + inner_client.start_workflow = AsyncMock() + + tc = TemporalClient(temporal_client=inner_client) + + with pytest.raises(ValueError, match="TERMINATE_EXISTING"): + await tc.start_workflow( + workflow="w", + arg={}, + id="id-1", + task_queue="q", + duplicate_policy=DuplicateWorkflowPolicy.TERMINATE_IF_RUNNING, + conflict_policy=ConflictWorkflowPolicy.USE_EXISTING, + ) + + inner_client.start_workflow.assert_not_awaited() + + +if __name__ == "__main__": # pragma: no cover + raise SystemExit(pytest.main([__file__, "-v"])) 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/lib/core/tracing/test_span_error.py b/tests/lib/core/tracing/test_span_error.py index 18116dcb3..02e9645a4 100644 --- a/tests/lib/core/tracing/test_span_error.py +++ b/tests/lib/core/tracing/test_span_error.py @@ -6,11 +6,19 @@ from unittest.mock import MagicMock, patch import pytest +from scale_gp_beta.lib.tracing import ( + PlatformError as SGPPlatformError, + ApplicationError as SGPApplicationError, + CategorizedError as SGPCategorizedError, +) from agentex.types.span import Span from agentex.lib.core.tracing.trace import Trace, AsyncTrace from agentex.lib.core.tracing.span_error import ( SPAN_ERROR_KEY, + PlatformError, + ApplicationError, + CategorizedError, get_span_error, set_span_error, ) @@ -34,12 +42,52 @@ def _make_span(data=None) -> Span: class TestSpanErrorHelpers: + def test_uses_canonical_sgp_error_types(self): + assert CategorizedError is SGPCategorizedError + assert ApplicationError is SGPApplicationError + assert PlatformError is SGPPlatformError + def test_set_then_get_on_none_data(self): span = _make_span(data=None) set_span_error(span, ValueError("boom")) - assert get_span_error(span) == {"type": "ValueError", "message": "boom"} + assert get_span_error(span) == { + "type": "ValueError", + "message": "boom", + "category": "unknown", + } assert isinstance(span.data, dict) - assert span.data[SPAN_ERROR_KEY] == {"type": "ValueError", "message": "boom"} + assert span.data[SPAN_ERROR_KEY] == { + "type": "ValueError", + "message": "boom", + "category": "unknown", + } + + def test_set_uses_explicit_exception_category(self): + span = _make_span(data=None) + set_span_error(span, PlatformError("unavailable")) + assert get_span_error(span) == { + "type": "PlatformError", + "message": "unavailable", + "category": "platform", + } + + def test_explicit_category_takes_precedence(self): + span = _make_span(data=None) + set_span_error(span, PlatformError("bad input"), error_category="application") + assert get_span_error(span)["category"] == "application" # type: ignore[index] + + def test_set_uses_application_error_category(self): + span = _make_span(data=None) + set_span_error(span, ApplicationError("bad input")) + assert get_span_error(span)["category"] == "application" # type: ignore[index] + + def test_bare_exception_attribute_does_not_opt_in(self): + class ImplicitlyCategorizedError(RuntimeError): + error_category = "platform" + + span = _make_span(data=None) + set_span_error(span, ImplicitlyCategorizedError("boom")) + assert get_span_error(span)["category"] == "unknown" # type: ignore[index] def test_set_preserves_existing_dict_keys(self): span = _make_span(data={"__span_type__": "LLM"}) @@ -76,7 +124,11 @@ def test_sync_span_records_error_and_reraises(self): captured["span"] = span raise ValueError("boom") err = get_span_error(captured["span"]) - assert err == {"type": "ValueError", "message": "boom"} + assert err == { + "type": "ValueError", + "message": "boom", + "category": "unknown", + } def test_sync_span_success_has_no_error(self): trace = Trace(processors=[], client=MagicMock(), trace_id="t1") @@ -93,7 +145,11 @@ async def test_async_span_records_error_and_reraises(self): captured["span"] = span raise RuntimeError("kaboom") err = get_span_error(captured["span"]) - assert err == {"type": "RuntimeError", "message": "kaboom"} + assert err == { + "type": "RuntimeError", + "message": "kaboom", + "category": "unknown", + } # --------------------------------------------------------------------------- @@ -111,7 +167,7 @@ def set_error( self, error_type: str | None = None, error_message: str | None = None, - exception: BaseException | None = None, + exception: BaseException | None = None, # noqa: ARG002 ) -> None: self.status = "ERROR" self.metadata["error"] = True @@ -131,7 +187,15 @@ def _env(): def test_error_maps_to_status_error(self): from agentex.lib.core.tracing.processors.sgp_tracing_processor import _build_sgp_span - span = _make_span(data={SPAN_ERROR_KEY: {"type": "ValueError", "message": "boom"}}) + span = _make_span( + data={ + SPAN_ERROR_KEY: { + "type": "ValueError", + "message": "boom", + "category": "application", + } + } + ) with patch(f"{PROCESSOR_MODULE}.create_span", side_effect=_fake_create_span): sgp_span = _build_sgp_span(span, self._env()) @@ -139,6 +203,7 @@ def test_error_maps_to_status_error(self): assert sgp_span.metadata["error"] is True assert sgp_span.metadata["error_type"] == "ValueError" assert sgp_span.metadata["error_message"] == "boom" + assert sgp_span.metadata["error_category"] == "application" def test_no_error_leaves_status_success(self): from agentex.lib.core.tracing.processors.sgp_tracing_processor import _build_sgp_span 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() == [] diff --git a/tests/test_adk_tracing_span_error.py b/tests/test_adk_tracing_span_error.py new file mode 100644 index 000000000..c81015142 --- /dev/null +++ b/tests/test_adk_tracing_span_error.py @@ -0,0 +1,120 @@ +"""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", + "category": "unknown", + } + + # 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", + "category": "unknown", + } + + +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", + "category": "unknown", + } + 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 == {} diff --git a/uv.lock b/uv.lock index f925c2dce..6c79c4ace 100644 --- a/uv.lock +++ b/uv.lock @@ -155,7 +155,7 @@ requires-dist = [ { name = "redis", specifier = ">=5.2.0,<8" }, { name = "rich", specifier = ">=13.9.2,<14" }, { name = "scale-gp", specifier = ">=0.1.0a59" }, - { name = "scale-gp-beta", specifier = ">=0.2.0" }, + { name = "scale-gp-beta", specifier = ">=0.5.0" }, { name = "starlette", specifier = ">=0.49.1" }, { name = "temporalio", specifier = ">=1.26.0,<2" }, { name = "typer", specifier = ">=0.16,<0.17" }, @@ -2791,7 +2791,7 @@ wheels = [ [[package]] name = "scale-gp-beta" -version = "0.2.0" +version = "0.5.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, @@ -2801,9 +2801,9 @@ dependencies = [ { name = "sniffio" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/3e/13/181f6b5a0e3fe5c2ca8b7b39e024ed11feba2cf0c879a6d77d84c8060383/scale_gp_beta-0.2.0.tar.gz", hash = "sha256:d4eac4a178ea4b7f21cfe2b421107009b69c4a5c1cbd2c7c864142c30ffda01c", size = 434620, upload-time = "2026-05-04T16:35:53.879Z" } +sdist = { url = "https://files.pythonhosted.org/packages/da/a6/623e122dd271f3c01852a65d57343bafbfd4b712068f63d49144c6210faa/scale_gp_beta-0.5.0.tar.gz", hash = "sha256:9f0de217d7bacd1880a7b9df6cf4f8be5d0620e24c382da14f7c0bd55423977e", size = 480740, upload-time = "2026-08-05T20:57:09.495Z" } 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" }, + { url = "https://files.pythonhosted.org/packages/50/95/4580c7e8d5e6d2354b09eb010d666f6359946a177b32dee8511d50a35f73/scale_gp_beta-0.5.0-py3-none-any.whl", hash = "sha256:1b1c6415a2c476c47ce5658bc0d8d2487b9aa61d95adb1002e0d5fbbe0b256c2", size = 472607, upload-time = "2026-08-05T20:57:07.9Z" }, ] [[package]]