feat: propagate error categories to SGP spans (#486)

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: jshaikScale

adk/pyproject.toml

@@ -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",

src/agentex/lib/adk/_modules/tracing.py

@@ -19,8 +19,8 @@
     StartSpanParams,
     TracingActivityName,
 )
-from agentex.lib.core.tracing.tracer import AsyncTracer
 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
 from agentex.lib.utils.logging import make_logger

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",

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
 
 

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):

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()
 

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,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

tests/test_adk_tracing_span_error.py

@@ -48,13 +48,21 @@ async def test_span_records_error_and_reraises() -> None:
             raise ValueError("boom")
 
     error = get_span_error(span)
-    assert error == {"type": "ValueError", "message": "boom"}
+    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"}
+    assert get_span_error(persisted_span) == {
+        "type": "ValueError",
+        "message": "boom",
+        "category": "unknown",
+    }
 
 
 async def test_span_success_records_no_error() -> None:
@@ -104,5 +112,9 @@ async def test_turn_span_records_error_and_reraises() -> None:
             assert turn.span is span
             raise ValueError("boom")
 
-    assert get_span_error(span) == {"type": "ValueError", "message": "boom"}
+    assert get_span_error(span) == {
+        "type": "ValueError",
+        "message": "boom",
+        "category": "unknown",
+    }
     end_span.assert_awaited_once()