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
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
from agentex.lib.utils.logging import make_logger
from agentex.lib.core.observability import tracing_metrics_recording as _metrics
from agentex.lib.environment_variables import EnvironmentVariables
from agentex.lib.core.tracing.span_error import get_span_error
from agentex.lib.core.tracing.processors.tracing_processor_interface import (
SyncTracingProcessor,
AsyncTracingProcessor,
Expand Down Expand Up @@ -83,6 +84,9 @@ def _build_sgp_span(span: Span, env_vars: EnvironmentVariables) -> SGPSpan:
),
)
sgp_span.start_time = span.start_time.isoformat() # type: ignore[union-attr]
error = get_span_error(span)
if error is not None:
sgp_span.set_error(error_type=error["type"], error_message=error["message"])
return sgp_span


Expand Down
36 changes: 36 additions & 0 deletions src/agentex/lib/core/tracing/span_error.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
from __future__ import annotations

from typing import Any

from agentex.types.span import Span

# Reserved key under ``Span.data`` carrying failure info for a span whose
# context-manager body raised. Mirrors the existing ``__span_type__`` /
# ``__source__`` reserved-key convention already read/written by the SGP
# processor. Stored in ``data`` because the Span model is generated from the
# OpenAPI spec and has no first-class status/error field; ``data`` is a real
# field, so it survives ``model_copy(deep=True)`` and round-trips to both the
# SGP and agentex-native span stores.
SPAN_ERROR_KEY = "__error__"


def set_span_error(span: Span, exc: BaseException) -> None:
"""Record an exception on ``span`` under ``data[SPAN_ERROR_KEY]``.

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)}
if span.data is None:
span.data = {}
if isinstance(span.data, dict):
span.data[SPAN_ERROR_KEY] = error


def get_span_error(span: Span) -> dict[str, Any] | None:
"""Return the error recorded by :func:`set_span_error`, or ``None``."""
if isinstance(span.data, dict):
value = span.data.get(SPAN_ERROR_KEY)
if isinstance(value, dict):
return value
return None
7 changes: 7 additions & 0 deletions src/agentex/lib/core/tracing/trace.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from agentex.types.span import Span
from agentex.lib.utils.logging import make_logger
from agentex.lib.utils.model_utils import recursive_model_dump
from agentex.lib.core.tracing.span_error import set_span_error
from agentex.lib.core.tracing.span_queue import (
SpanEventType,
AsyncSpanQueue,
Expand Down Expand Up @@ -165,6 +166,9 @@ def span(
span = self.start_span(name, parent_id, input, data, task_id=task_id)
try:
yield span
except Exception as exc:
set_span_error(span, exc)
raise
finally:
self.end_span(span)

Expand Down Expand Up @@ -321,5 +325,8 @@ async def span(
span = await self.start_span(name, parent_id, input, data, task_id=task_id)
try:
yield span
except Exception as exc:
set_span_error(span, exc)
raise
finally:
await self.end_span(span)
151 changes: 151 additions & 0 deletions tests/lib/core/tracing/test_span_error.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,151 @@
from __future__ import annotations

import uuid
from typing import Any
from datetime import UTC, datetime
from unittest.mock import MagicMock, patch

import pytest

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,
get_span_error,
set_span_error,
)

PROCESSOR_MODULE = "agentex.lib.core.tracing.processors.sgp_tracing_processor"


def _make_span(data=None) -> Span:
return Span(
id=str(uuid.uuid4()),
name="test-span",
start_time=datetime.now(UTC),
trace_id="trace-1",
data=data,
)


# ---------------------------------------------------------------------------
# Helpers: set_span_error / get_span_error
# ---------------------------------------------------------------------------


class TestSpanErrorHelpers:
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 isinstance(span.data, dict)
assert span.data[SPAN_ERROR_KEY] == {"type": "ValueError", "message": "boom"}

def test_set_preserves_existing_dict_keys(self):
span = _make_span(data={"__span_type__": "LLM"})
set_span_error(span, RuntimeError("nope"))
assert isinstance(span.data, dict)
assert span.data["__span_type__"] == "LLM"
err = get_span_error(span)
assert err is not None
assert err["type"] == "RuntimeError"

def test_get_returns_none_when_no_error(self):
assert get_span_error(_make_span(data={"foo": "bar"})) is None
assert get_span_error(_make_span(data=None)) is None

def test_set_is_noop_on_list_data(self):
span = _make_span(data=[{"a": 1}])
set_span_error(span, ValueError("boom"))
# list-shaped data is left untouched (mirrors _add_source_to_span)
assert span.data == [{"a": 1}]
assert get_span_error(span) is None


# ---------------------------------------------------------------------------
# Capture: the context managers record body exceptions onto the span
# ---------------------------------------------------------------------------


class TestContextManagerCapture:
def test_sync_span_records_error_and_reraises(self):
trace = Trace(processors=[], client=MagicMock(), trace_id="t1")
captured = {}
with pytest.raises(ValueError, match="boom"):
with trace.span("op") as span:
captured["span"] = span
raise ValueError("boom")
err = get_span_error(captured["span"])
assert err == {"type": "ValueError", "message": "boom"}

def test_sync_span_success_has_no_error(self):
trace = Trace(processors=[], client=MagicMock(), trace_id="t1")
with trace.span("op") as span:
pass
assert get_span_error(span) is None

@pytest.mark.asyncio
async def test_async_span_records_error_and_reraises(self):
trace = AsyncTrace(processors=[], client=MagicMock(), trace_id="t1")
captured = {}
with pytest.raises(RuntimeError, match="kaboom"):
async with trace.span("op") as span:
captured["span"] = span
raise RuntimeError("kaboom")
err = get_span_error(captured["span"])
assert err == {"type": "RuntimeError", "message": "kaboom"}


# ---------------------------------------------------------------------------
# Map: _build_sgp_span translates the recorded error into SGP status=ERROR
# ---------------------------------------------------------------------------


class _FakeSGPSpan:
def __init__(self, metadata: dict[str, Any] | None) -> None:
self.status = "SUCCESS"
self.metadata: dict[str, Any] = metadata if metadata is not None else {}
self.start_time = None

def set_error(
self,
error_type: str | None = None,
error_message: str | None = None,
exception: BaseException | None = None,
) -> None:
self.status = "ERROR"
self.metadata["error"] = True
self.metadata["error_type"] = error_type
self.metadata["error_message"] = error_message


def _fake_create_span(**kwargs: Any) -> _FakeSGPSpan:
return _FakeSGPSpan(kwargs.get("metadata"))


class TestBuildSGPSpanMapping:
@staticmethod
def _env():
return MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None)

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"}})
with patch(f"{PROCESSOR_MODULE}.create_span", side_effect=_fake_create_span):
sgp_span = _build_sgp_span(span, self._env())

assert sgp_span.status == "ERROR"
assert sgp_span.metadata["error"] is True
assert sgp_span.metadata["error_type"] == "ValueError"
assert sgp_span.metadata["error_message"] == "boom"

def test_no_error_leaves_status_success(self):
from agentex.lib.core.tracing.processors.sgp_tracing_processor import _build_sgp_span

span = _make_span(data={"__span_type__": "LLM"})
with patch(f"{PROCESSOR_MODULE}.create_span", side_effect=_fake_create_span):
sgp_span = _build_sgp_span(span, self._env())

assert sgp_span.status == "SUCCESS"
assert "error" not in sgp_span.metadata
Loading