Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions src/agentex/lib/adk/_modules/tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
TracingActivityName,
)
from agentex.lib.core.tracing.tracer import AsyncTracer
from agentex.lib.core.tracing.span_error import set_span_error
from agentex.lib.core.harness.types import TurnUsage
from agentex.types.span import Span
from agentex.lib.utils.logging import make_logger
Expand Down Expand Up @@ -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(
Expand Down
103 changes: 76 additions & 27 deletions src/agentex/lib/core/temporal/services/temporal_task_service.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,10 @@
from __future__ import annotations

import sys
from typing import Any
from datetime import timedelta
from contextlib import contextmanager
from collections.abc import Iterator

from agentex.types.task import Task
from agentex.types.agent import Agent
Expand All @@ -13,6 +16,55 @@
from agentex.lib.core.clients.temporal.temporal_client import TemporalClient


@contextmanager
def _acp_dispatch_span(name: str, task_id: str | None = None) -> Iterator[None]:
"""Wrap an ACP -> Temporal dispatch (start_workflow / signal) in an OTel span.

The Temporal OpenTelemetry interceptor propagates trace context by injecting
the CURRENTLY ACTIVE span into the Temporal message headers on the caller
side (``start_workflow`` / ``signal_workflow``); the worker then extracts it
and roots the workflow / activity spans under it. But the ACP server dispatches
from a bare async handler with no active span, so nothing is injected and the
workflow's activities become DETACHED trace roots -- the business work shows up
in Tempo as a fresh trace with no link back to the ``task/create`` /
``event/send`` that triggered it.

Opening a span here gives the interceptor something to inject. It becomes a
child of the ingress request span when one is active (front-of-request
propagation), or a fresh per-turn root otherwise.

Fail-open across the WHOLE obs setup, not just the import: ``get_tracer`` and
entering ``start_as_current_span`` run the sampler and every
``SpanProcessor.on_start`` (the SDK does not guard those), so a broken
provider or a custom sampler/processor that raises would otherwise fail the
dispatch itself. If any of it fails we run the dispatch untraced. The dispatch
body (the ``yield``) is OUTSIDE the guard so its exceptions still propagate.
"""
span_cm = None
try:
from opentelemetry import trace as _otel_trace

tracer = _otel_trace.get_tracer("agentex.acp")
# task_id goes on an attribute, NOT in the span name: a per-task span name is
# high-cardinality and breaks span-name aggregation in Tempo.
attributes = {"agentex.task_id": task_id} if task_id else None
span_cm = tracer.start_as_current_span(name, kind=_otel_trace.SpanKind.PRODUCER, attributes=attributes)
span_cm.__enter__()
except Exception: # pragma: no cover - obs must never break a dispatch
span_cm = None

try:
yield
finally:
if span_cm is not None:
# Pass exc info so the span reflects a failed dispatch; guard __exit__
# so closing the span can never mask the dispatch outcome.
try:
span_cm.__exit__(*sys.exc_info())
except Exception: # pragma: no cover - best-effort close
pass


class TemporalTaskService:
"""
Submits Agent agent_tasks to the async runtime for execution.
Expand All @@ -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.
Expand All @@ -37,22 +88,19 @@ async def submit_task(self, agent: Agent, task: Task, params: dict[str, Any] | N
# indefinitely, which long-lived chat/session agents rely on). A positive
# value bounds the whole continue-as-new chain's wall-clock lifetime.
timeout_seconds = self._env_vars.WORKFLOW_EXECUTION_TIMEOUT_SECONDS
execution_timeout = (
timedelta(seconds=timeout_seconds)
if timeout_seconds and timeout_seconds > 0
else None
)
return await self._temporal_client.start_workflow(
workflow=self._env_vars.WORKFLOW_NAME,
arg=CreateTaskParams(
agent=agent,
task=task,
params=params,
),
id=task.id,
task_queue=self._env_vars.WORKFLOW_TASK_QUEUE,
execution_timeout=execution_timeout,
)
execution_timeout = timedelta(seconds=timeout_seconds) if timeout_seconds and timeout_seconds > 0 else None
with _acp_dispatch_span("acp.task_create", task_id=task.id):
return await self._temporal_client.start_workflow(
workflow=self._env_vars.WORKFLOW_NAME,
arg=CreateTaskParams(
agent=agent,
task=task,
params=params,
),
id=task.id,
task_queue=self._env_vars.WORKFLOW_TASK_QUEUE,
execution_timeout=execution_timeout,
)

async def get_state(self, task_id: str) -> WorkflowState:
"""
Expand All @@ -63,16 +111,17 @@ async def get_state(self, task_id: str) -> WorkflowState:
)

async def send_event(self, agent: Agent, task: Task, event: Event, request: dict | None = None) -> None:
return await self._temporal_client.send_signal(
workflow_id=task.id,
signal=SignalName.RECEIVE_EVENT.value,
payload=SendEventParams(
agent=agent,
task=task,
event=event,
request=request,
).model_dump(),
)
with _acp_dispatch_span("acp.event_send", task_id=task.id):
return await self._temporal_client.send_signal(
workflow_id=task.id,
signal=SignalName.RECEIVE_EVENT.value,
payload=SendEventParams(
agent=agent,
task=task,
event=event,
request=request,
).model_dump(),
)

async def interrupt(self, agent: Agent, task: Task, request: dict | None = None) -> None:
"""Forward a task/interrupt to the running workflow as a dedicated signal.
Expand Down
44 changes: 30 additions & 14 deletions src/agentex/lib/core/tracing/obs_ids.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand Down Expand Up @@ -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]}
Loading