diff --git a/src/sentry/spans/consumers/process_segments/enrichment.py b/src/sentry/spans/consumers/process_segments/enrichment.py index b2736cb92a6d9d..fda3ee1819d8a2 100644 --- a/src/sentry/spans/consumers/process_segments/enrichment.py +++ b/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,6 +122,20 @@ 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), @@ -126,6 +143,24 @@ def get_value(key: str) -> Any: 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])) diff --git a/tests/sentry/spans/consumers/process_segments/test_enrichment.py b/tests/sentry/spans/consumers/process_segments/test_enrichment.py index 6c9ec9071d7a04..494c4adf8b36a2 100644 --- a/tests/sentry/spans/consumers/process_segments/test_enrichment.py +++ b/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