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

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
@@ -76,7 +76,10 @@ def __init__(self, spans: list[SpanEvent]) -> None:
         self._ttfd_ts = _timestamp_by_op(spans, "ui.load.full_display")
 
         self._span_map: dict[str, list[tuple[int, int]]] = {}
+        self._span_hierarchy: dict[str, SpanEvent] = {}
         for span in spans:
+            if "span_id" in span:
+                self._span_hierarchy[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)
@@ -119,13 +122,34 @@ 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:
+                for ancestor in self._iter_ancestors(span):
+                    if (
+                        get_span_op(ancestor) == "gen_ai.invoke_agent"
+                        and (agent_name := attribute_value(ancestor, "gen_ai.agent.name"))
+                        is not None
+                    ):
+                        attributes["gen_ai.agent.name"] = {
+                            "type": "string",
+                            "value": agent_name,
+                        }
+                        break
+
         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 = span
+        while current := self._span_hierarchy.get(current.get("parent_span_id")):
+            yield current
+
     def _exclusive_time(self, span: SpanEvent) -> float:
         """
         Sets the exclusive time on all spans in the list.

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

@@ -471,3 +471,183 @@ 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."""
+    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": "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_from_ancestor() -> None:
+    """Test that gen_ai.agent.name is inherited from a grandparent ancestor."""
+    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.create_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") == "GrandparentAgent"
+
+
+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.create_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.create_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