feat(consumers): ensure that gen_ai spans get the gen_ai.agent.name from the ancestor span (#102955)

Author: constantinius

src/sentry/spans/consumers/process_segments/enrichment.py

@@ -1,5 +1,5 @@
 from collections import defaultdict
-from collections.abc import Sequence
+from collections.abc import Iterator, Sequence
 from typing import Any
 
 from sentry_kafka_schemas.schema_types.ingest_spans_v1 import SpanEvent
@@ -75,11 +75,14 @@ def __init__(self, spans: list[SpanEvent]) -> None:
         self._ttid_ts = _timestamp_by_op(spans, "ui.load.initial_display")
         self._ttfd_ts = _timestamp_by_op(spans, "ui.load.full_display")
 
-        self._span_map: dict[str, list[tuple[int, int]]] = {}
+        self._span_intervals: dict[str, list[tuple[int, int]]] = {}
+        self._spans_by_id: dict[str, SpanEvent] = {}
         for span in spans:
+            if "span_id" in span:
+                self._spans_by_id[span["span_id"]] = span
             if parent_span_id := span.get("parent_span_id"):
                 interval = _span_interval(span)
-                self._span_map.setdefault(parent_span_id, []).append(interval)
+                self._span_intervals.setdefault(parent_span_id, []).append(interval)
 
     def _attributes(self, span: SpanEvent) -> dict[str, Any]:
         attributes: dict[str, Any] = {**(span.get("attributes") or {})}
@@ -119,13 +122,45 @@ def get_value(key: str) -> Any:
                 if attributes.get(key) is None:
                     attributes[key] = value
 
+            if get_span_op(span).startswith("gen_ai.") and "gen_ai.agent.name" not in attributes:
+                if (parent_span_id := span.get("parent_span_id")) is not None:
+                    parent_span = self._spans_by_id.get(parent_span_id)
+                    if (
+                        parent_span is not None
+                        and get_span_op(parent_span) == "gen_ai.invoke_agent"
+                        and (agent_name := attribute_value(parent_span, "gen_ai.agent.name"))
+                        is not None
+                    ):
+                        attributes["gen_ai.agent.name"] = {
+                            "type": "string",
+                            "value": agent_name,
+                        }
+
         attributes["sentry.exclusive_time_ms"] = {
             "type": "double",
             "value": self._exclusive_time(span),
         }
 
         return attributes
 
+    def _iter_ancestors(self, span: SpanEvent) -> Iterator[SpanEvent]:
+        """
+        Iterates over the ancestors of a span in order towards the root using the "parent_span_id" attribute.
+        """
+        current: SpanEvent | None = span
+        parent_span_id: str | None = None
+
+        while current is not None:
+            parent_span_id = current.get("parent_span_id")
+            if parent_span_id is not None:
+                current = self._spans_by_id.get(parent_span_id)
+            else:
+                current = None
+            if current is not None:
+                yield current
+            else:
+                break
+
     def _exclusive_time(self, span: SpanEvent) -> float:
         """
         Sets the exclusive time on all spans in the list.
@@ -134,7 +169,7 @@ def _exclusive_time(self, span: SpanEvent) -> float:
         of all time intervals where no child span was active.
         """
 
-        intervals = self._span_map.get(span["span_id"], [])
+        intervals = self._span_intervals.get(span["span_id"], [])
         # Sort by start ASC, end DESC to skip over nested intervals efficiently
         intervals.sort(key=lambda x: (x[0], -x[1]))
 

tests/sentry/spans/consumers/process_segments/test_enrichment.py

@@ -471,3 +471,215 @@ def _mock_performance_issue_span(is_segment, attributes, **fields) -> SpanEvent:
             **fields,
         },
     )
+
+
+def test_enrich_gen_ai_agent_name_from_immediate_parent() -> None:
+    """Test that gen_ai.agent.name is inherited from the immediate parent with gen_ai.invoke_agent operation."""
+    parent_span = build_mock_span(
+        project_id=1,
+        is_segment=True,
+        span_id="aaaaaaaaaaaaaaaa",
+        start_timestamp=1609455600.0,
+        end_timestamp=1609455605.0,
+        span_op="gen_ai.invoke_agent",
+        attributes={
+            "gen_ai.agent.name": {"type": "string", "value": "MyAgent"},
+        },
+    )
+
+    child_span = build_mock_span(
+        project_id=1,
+        span_id="bbbbbbbbbbbbbbbb",
+        parent_span_id="aaaaaaaaaaaaaaaa",
+        start_timestamp=1609455601.0,
+        end_timestamp=1609455602.0,
+        span_op="gen_ai.execute_tool",
+    )
+
+    spans = [parent_span, child_span]
+    _, enriched_spans = TreeEnricher.enrich_spans(spans)
+    compatible_spans = [make_compatible(span) for span in enriched_spans]
+
+    parent, child = compatible_spans
+    assert attribute_value(parent, "gen_ai.agent.name") == "MyAgent"
+    assert attribute_value(child, "gen_ai.agent.name") == "MyAgent"
+
+
+def test_enrich_gen_ai_agent_name_not_overwritten() -> None:
+    """Test that gen_ai.agent.name is not overwritten if already set on the child."""
+    parent_span = build_mock_span(
+        project_id=1,
+        is_segment=True,
+        span_id="aaaaaaaaaaaaaaaa",
+        start_timestamp=1609455600.0,
+        end_timestamp=1609455605.0,
+        span_op="gen_ai.invoke_agent",
+        attributes={
+            "gen_ai.agent.name": {"type": "string", "value": "ParentAgent"},
+        },
+    )
+
+    child_span = build_mock_span(
+        project_id=1,
+        span_id="bbbbbbbbbbbbbbbb",
+        parent_span_id="aaaaaaaaaaaaaaaa",
+        start_timestamp=1609455601.0,
+        end_timestamp=1609455602.0,
+        span_op="gen_ai.handoff",
+        attributes={
+            "gen_ai.agent.name": {"type": "string", "value": "ChildAgent"},
+        },
+    )
+
+    spans = [parent_span, child_span]
+    _, enriched_spans = TreeEnricher.enrich_spans(spans)
+    compatible_spans = [make_compatible(span) for span in enriched_spans]
+
+    parent, child = compatible_spans
+    assert attribute_value(parent, "gen_ai.agent.name") == "ParentAgent"
+    assert attribute_value(child, "gen_ai.agent.name") == "ChildAgent"
+
+
+def test_enrich_gen_ai_agent_name_not_set_without_ancestor() -> None:
+    """Test that gen_ai.agent.name is not set if no ancestor has it."""
+    parent_span = build_mock_span(
+        project_id=1,
+        is_segment=True,
+        span_id="aaaaaaaaaaaaaaaa",
+        start_timestamp=1609455600.0,
+        end_timestamp=1609455605.0,
+        span_op="some.operation",
+    )
+
+    child_span = build_mock_span(
+        project_id=1,
+        span_id="bbbbbbbbbbbbbbbb",
+        parent_span_id="aaaaaaaaaaaaaaaa",
+        start_timestamp=1609455601.0,
+        end_timestamp=1609455602.0,
+        span_op="gen_ai.execute_tool",
+    )
+
+    spans = [parent_span, child_span]
+    _, enriched_spans = TreeEnricher.enrich_spans(spans)
+    compatible_spans = [make_compatible(span) for span in enriched_spans]
+
+    parent, child = compatible_spans
+    assert attribute_value(parent, "gen_ai.agent.name") is None
+    assert attribute_value(child, "gen_ai.agent.name") is None
+
+
+def test_enrich_gen_ai_agent_name_not_from_sibling() -> None:
+    """Test that gen_ai.agent.name is not taken from a sibling span."""
+    parent_span = build_mock_span(
+        project_id=1,
+        is_segment=True,
+        span_id="aaaaaaaaaaaaaaaa",
+        start_timestamp=1609455600.0,
+        end_timestamp=1609455605.0,
+        span_op="some.operation",
+    )
+
+    sibling_with_agent = build_mock_span(
+        project_id=1,
+        span_id="bbbbbbbbbbbbbbbb",
+        parent_span_id="aaaaaaaaaaaaaaaa",
+        start_timestamp=1609455601.0,
+        end_timestamp=1609455602.0,
+        span_op="gen_ai.invoke_agent",
+        attributes={
+            "gen_ai.agent.name": {"type": "string", "value": "SiblingAgent"},
+        },
+    )
+
+    target_child = build_mock_span(
+        project_id=1,
+        span_id="cccccccccccccccc",
+        parent_span_id="aaaaaaaaaaaaaaaa",
+        start_timestamp=1609455602.5,
+        end_timestamp=1609455603.5,
+        span_op="gen_ai.execute_tool",
+    )
+
+    spans = [parent_span, sibling_with_agent, target_child]
+    _, enriched_spans = TreeEnricher.enrich_spans(spans)
+    compatible_spans = [make_compatible(span) for span in enriched_spans]
+
+    parent, sibling, target = compatible_spans
+    assert attribute_value(parent, "gen_ai.agent.name") is None
+    assert attribute_value(sibling, "gen_ai.agent.name") == "SiblingAgent"
+    assert attribute_value(target, "gen_ai.agent.name") is None
+
+
+def test_enrich_gen_ai_agent_name_only_from_invoke_agent_parent() -> None:
+    """Test that gen_ai.agent.name is only inherited from parent with gen_ai.invoke_agent operation."""
+    parent_span = build_mock_span(
+        project_id=1,
+        is_segment=True,
+        span_id="aaaaaaaaaaaaaaaa",
+        start_timestamp=1609455600.0,
+        end_timestamp=1609455605.0,
+        span_op="gen_ai.create_agent",
+        attributes={
+            "gen_ai.agent.name": {"type": "string", "value": "CreateAgentName"},
+        },
+    )
+
+    child_span = build_mock_span(
+        project_id=1,
+        span_id="bbbbbbbbbbbbbbbb",
+        parent_span_id="aaaaaaaaaaaaaaaa",
+        start_timestamp=1609455601.0,
+        end_timestamp=1609455602.0,
+        span_op="gen_ai.run",
+    )
+
+    spans = [parent_span, child_span]
+    _, enriched_spans = TreeEnricher.enrich_spans(spans)
+    compatible_spans = [make_compatible(span) for span in enriched_spans]
+
+    parent, child = compatible_spans
+    assert attribute_value(parent, "gen_ai.agent.name") == "CreateAgentName"
+    assert attribute_value(child, "gen_ai.agent.name") is None
+
+
+def test_enrich_gen_ai_agent_name_not_from_grandparent() -> None:
+    """Test that gen_ai.agent.name is NOT inherited from grandparent, only from immediate parent."""
+    grandparent_span = build_mock_span(
+        project_id=1,
+        is_segment=True,
+        span_id="aaaaaaaaaaaaaaaa",
+        start_timestamp=1609455600.0,
+        end_timestamp=1609455605.0,
+        span_op="gen_ai.invoke_agent",
+        attributes={
+            "gen_ai.agent.name": {"type": "string", "value": "GrandparentAgent"},
+        },
+    )
+
+    parent_span = build_mock_span(
+        project_id=1,
+        span_id="bbbbbbbbbbbbbbbb",
+        parent_span_id="aaaaaaaaaaaaaaaa",
+        start_timestamp=1609455601.0,
+        end_timestamp=1609455604.0,
+        span_op="some.operation",
+    )
+
+    child_span = build_mock_span(
+        project_id=1,
+        span_id="cccccccccccccccc",
+        parent_span_id="bbbbbbbbbbbbbbbb",
+        start_timestamp=1609455602.0,
+        end_timestamp=1609455603.0,
+        span_op="gen_ai.run",
+    )
+
+    spans = [grandparent_span, parent_span, child_span]
+    _, enriched_spans = TreeEnricher.enrich_spans(spans)
+    compatible_spans = [make_compatible(span) for span in enriched_spans]
+
+    grandparent, parent, child = compatible_spans
+    assert attribute_value(grandparent, "gen_ai.agent.name") == "GrandparentAgent"
+    assert attribute_value(parent, "gen_ai.agent.name") is None
+    assert attribute_value(child, "gen_ai.agent.name") is None