Merge branch 'master' into webb/langchain/watched-span

Author: alexander-alderman-webb

sentry_sdk/integrations/huey.py

@@ -34,6 +34,13 @@
 except ImportError:
     raise DidNotEnable("Huey is not installed")
 
+try:
+    from huey.api import chord as HueyChord
+    from huey.api import group as HueyGroup
+except ImportError:
+    HueyChord = None
+    HueyGroup = None
+
 
 HUEY_CONTROL_FLOW_EXCEPTIONS = (CancelExecution, RetryTask, TaskLockedException)
 
@@ -53,22 +60,37 @@ def patch_enqueue() -> None:
 
     @ensure_integration_enabled(HueyIntegration, old_enqueue)
     def _sentry_enqueue(
-        self: "Huey", task: "Task"
+        self: "Huey", item: "Union[Task, HueyGroup, HueyChord]"
     ) -> "Optional[Union[Result, ResultGroup]]":
+        if HueyChord is not None and isinstance(item, HueyChord):
+            span_name = "Huey Chord"
+        elif HueyGroup is not None and isinstance(item, HueyGroup):
+            span_name = "Huey Task Group"
+        else:
+            span_name = item.name
+
         with sentry_sdk.start_span(
             op=OP.QUEUE_SUBMIT_HUEY,
-            name=task.name,
+            name=span_name,
             origin=HueyIntegration.origin,
         ):
-            if not isinstance(task, PeriodicTask):
+            if (
+                not isinstance(item, PeriodicTask)
+                and not (HueyGroup is not None and isinstance(item, HueyGroup))
+                and not (HueyChord is not None and isinstance(item, HueyChord))
+            ):
                 # Attach trace propagation data to task kwargs. We do
                 # not do this for periodic tasks, as these don't
                 # really have an originating transaction.
-                task.kwargs["sentry_headers"] = {
+                # Additionally, we do not do this for Huey groups or chords, as enqueue will
+                # recursively call this method for each task within the list, resulting
+                # in the trace propagation data being attached to each task individually
+                # (which we want)
+                item.kwargs["sentry_headers"] = {
                     BAGGAGE_HEADER_NAME: get_baggage(),
                     SENTRY_TRACE_HEADER_NAME: get_traceparent(),
                 }
-            return old_enqueue(self, task)
+            return old_enqueue(self, item)
 
     Huey.enqueue = _sentry_enqueue
 

sentry_sdk/integrations/openai_agents/spans/ai_client.py

@@ -5,7 +5,6 @@
 
 from ..consts import SPAN_ORIGIN
 from ..utils import (
-    _create_mcp_execute_tool_spans,
     _set_agent_data,
     _set_input_data,
     _set_output_data,
@@ -55,7 +54,6 @@ def update_ai_client_span(
 
     if hasattr(response, "output") and response.output:
         _set_output_data(span, response)
-        _create_mcp_execute_tool_spans(span, response)
 
     if response_model is not None:
         span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response_model)

sentry_sdk/integrations/openai_agents/utils.py

@@ -14,7 +14,7 @@
     set_data_normalized,
     truncate_and_annotate_messages,
 )
-from sentry_sdk.consts import OP, SPANDATA, SPANSTATUS
+from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations import DidNotEnable
 from sentry_sdk.scope import should_send_default_pii
 from sentry_sdk.utils import event_from_exception, safe_serialize
@@ -215,25 +215,3 @@ def _set_output_data(span: "sentry_sdk.tracing.Span", result: "Any") -> None:
         set_data_normalized(
             span, SPANDATA.GEN_AI_RESPONSE_TEXT, output_messages["response"]
         )
-
-
-def _create_mcp_execute_tool_spans(
-    span: "sentry_sdk.tracing.Span", result: "agents.Result"
-) -> None:
-    for output in result.output:
-        if output.__class__.__name__ == "McpCall":
-            with sentry_sdk.start_span(
-                op=OP.GEN_AI_EXECUTE_TOOL,
-                name=f"execute_tool {output.name}",
-                start_timestamp=span.start_timestamp,
-            ) as execute_tool_span:
-                execute_tool_span.set_data(SPANDATA.GEN_AI_TOOL_NAME, output.name)
-                if should_send_default_pii():
-                    execute_tool_span.set_data(
-                        SPANDATA.GEN_AI_TOOL_INPUT, output.arguments
-                    )
-                    execute_tool_span.set_data(
-                        SPANDATA.GEN_AI_TOOL_OUTPUT, output.output
-                    )
-                if output.error:
-                    execute_tool_span.set_status(SPANSTATUS.INTERNAL_ERROR)

tests/integrations/huey/test_huey.py

@@ -11,6 +11,12 @@
 
 HUEY_VERSION = parse_version(HUEY_VERSION)
 
+try:
+    from huey.api import chord, group
+except ImportError:
+    chord = None
+    group = None
+
 
 @pytest.fixture
 def init_huey(sentry_init):
@@ -222,3 +228,114 @@ def propagated_trace_task():
     (event,) = events
 
     assert event["contexts"]["trace"]["origin"] == "auto.queue.huey"
+
+
+@pytest.mark.skipif(HUEY_VERSION < (3, 0), reason="group was added in 3.0")
+def test_huey_enqueue_group(init_huey, capture_events):
+    huey = init_huey()
+
+    events = capture_events()
+
+    @huey.task()
+    def task1():
+        pass
+
+    @huey.task()
+    def task2():
+        pass
+
+    with start_transaction() as transaction:
+        huey.enqueue(group([task1.s(), task2.s()]))
+
+    for _ in range(2):
+        task = huey.dequeue()
+        huey.execute(task)
+
+    assert len(events) == 3
+
+    # Assert enqueue spans were successfully recorded
+    producer_event = events[0]
+    assert producer_event["type"] == "transaction"
+    assert producer_event["contexts"]["trace"]["trace_id"] == transaction.trace_id
+    assert producer_event["contexts"]["trace"]["origin"] == "manual"
+
+    spans = producer_event["spans"]
+    assert len(spans) == 3
+    assert spans[0]["op"] == "queue.submit.huey"
+    assert spans[0]["description"] == "Huey Task Group"
+    assert spans[1]["op"] == "queue.submit.huey"
+    assert spans[1]["description"] == "task1"
+    assert spans[2]["op"] == "queue.submit.huey"
+    assert spans[2]["description"] == "task2"
+
+    # Consumer transaction assertions (one per task)
+    consumer_events = events[1:]
+    for _, (consumer_event, expected_name) in enumerate(
+        zip(consumer_events, ["task1", "task2"])
+    ):
+        assert consumer_event["type"] == "transaction"
+        assert consumer_event["transaction"] == expected_name
+        assert consumer_event["transaction_info"] == {"source": "task"}
+        assert consumer_event["contexts"]["trace"]["op"] == "queue.task.huey"
+        assert consumer_event["contexts"]["trace"]["origin"] == "auto.queue.huey"
+        assert consumer_event["contexts"]["trace"]["status"] == "ok"
+        assert consumer_event["contexts"]["trace"]["trace_id"] == transaction.trace_id
+        assert "huey_task_id" in consumer_event["tags"]
+        assert consumer_event["tags"]["huey_task_retry"] is False
+
+
+@pytest.mark.skipif(HUEY_VERSION < (3, 0), reason="chord was added in 3.0")
+def test_huey_enqueue_chord(init_huey, capture_events):
+    huey = init_huey()
+
+    events = capture_events()
+
+    @huey.task()
+    def task1():
+        pass
+
+    @huey.task()
+    def task2(results):
+        pass
+
+    with start_transaction() as transaction:
+        huey.enqueue(chord([task1.s()], task2.s()))
+
+    for _ in range(2):
+        task = huey.dequeue()
+        huey.execute(task)
+
+    assert len(events) == 3
+
+    # Enqueue spans
+    producer_event = events[0]
+    assert producer_event["contexts"]["trace"]["trace_id"] == transaction.trace_id
+    assert producer_event["contexts"]["trace"]["origin"] == "manual"
+
+    spans = producer_event["spans"]
+    assert len(spans) == 2
+    assert spans[0]["op"] == "queue.submit.huey"
+    assert spans[0]["description"] == "Huey Chord"
+    assert spans[1]["op"] == "queue.submit.huey"
+    assert spans[1]["description"] == "task1"
+
+    task1_event = events[1]
+    # Confirm the first task enqueued the chord callback
+    task1_spans = task1_event["spans"]
+    assert len(task1_spans) == 1
+    assert task1_spans[0]["op"] == "queue.submit.huey"
+    assert task1_spans[0]["description"] == "task2"
+
+    consumer_events = events[1:]
+    for _, (consumer_event, expected_name) in enumerate(
+        zip(consumer_events, ["task1", "task2"])
+    ):
+        assert consumer_event["type"] == "transaction"
+        assert consumer_event["transaction"] == expected_name
+        assert consumer_event["transaction_info"] == {"source": "task"}
+        assert consumer_event["contexts"]["trace"]["op"] == "queue.task.huey"
+        assert consumer_event["contexts"]["trace"]["origin"] == "auto.queue.huey"
+        assert consumer_event["contexts"]["trace"]["status"] == "ok"
+        assert consumer_event["contexts"]["trace"]["trace_id"] == transaction.trace_id
+        assert "huey_task_id" in consumer_event["tags"]
+        assert consumer_event["tags"]["huey_task_retry"] is False