Skip to content
Open
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
2 changes: 1 addition & 1 deletion adk/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
5 changes: 5 additions & 0 deletions src/agentex/lib/adk/_modules/tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -236,6 +237,10 @@ async def span(
)
try:
yield span
except Exception as exc:
if span:
set_span_error(span, exc)
raise
finally:
if span:
await self.end_span(
Expand Down
10 changes: 10 additions & 0 deletions src/agentex/lib/core/tracing/__init__.py
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -13,6 +19,10 @@
"Span",
"Tracer",
"AsyncTracer",
"CategorizedError",
"ApplicationError",
"PlatformError",
"ErrorCategory",
"AsyncSpanQueue",
"get_default_span_queue",
"shutdown_default_span_queue",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down
47 changes: 44 additions & 3 deletions src/agentex/lib/core/tracing/span_error.py
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -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):
Expand Down
19 changes: 19 additions & 0 deletions tests/lib/adk/test_tracing_module.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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()

Expand Down
77 changes: 71 additions & 6 deletions tests/lib/core/tracing/test_span_error.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand All @@ -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"})
Expand Down Expand Up @@ -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")
Expand All @@ -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",
}


# ---------------------------------------------------------------------------
Expand All @@ -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
Expand All @@ -131,14 +187,23 @@ 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())

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"
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
Expand Down
8 changes: 4 additions & 4 deletions uv.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading