fix(temporal): include activity retry attempts in metadata

py/src/braintrust/integrations/temporal/plugin.py

@@ -227,6 +227,7 @@ async def execute_activity(self, input: temporalio.worker.ExecuteActivityInput)
             metadata={
                 "temporal.activity_type": info.activity_type,
                 "temporal.activity_id": info.activity_id,
+                "temporal.activity_attempt": info.attempt,
                 "temporal.workflow_id": info.workflow_id,
                 "temporal.workflow_run_id": info.workflow_run_id,
             },

py/src/braintrust/integrations/temporal/test_temporal.py

@@ -205,22 +205,6 @@ async def run(self, input: TaskInput) -> int:
         return result
 
 
-@temporalio.workflow.defn
-class ParentWorkflow:
-    """Parent workflow that spawns a child workflow."""
-
-    @temporalio.workflow.run
-    async def run(self, input: TaskInput) -> int:
-        # Execute child workflow
-        child_result = await temporalio.workflow.execute_child_workflow(
-            ChildWorkflow.run,
-            input,
-            id=f"child-{temporalio.workflow.info().workflow_id}",
-        )
-
-        return child_result
-
-
 @temporalio.workflow.defn
 class DelegatingParentWorkflow:
     """Parent workflow that starts a child workflow and an activity inside an in-workflow span.
@@ -459,7 +443,7 @@ async def test_plugin_activity_after_replay_stays_under_workflow_span(
             assert workflow_span["root_span_id"] == client_spans[0]["root_span_id"]
 
     @pytest.mark.parametrize(
-        "workflow,input_value,expected_result,expected_span_counts,expected_activity_errors",
+        "workflow,input_value,expected_result,expected_span_counts,expected_activity_errors,expected_activity_attempts",
         [
             # failing_activity fails its first attempt, so the retry yields a second activity span.
             pytest.param(
@@ -468,27 +452,17 @@ async def test_plugin_activity_after_replay_stays_under_workflow_span(
                 50,
                 {"temporal.workflow.WorkflowWithRetry": 1, "temporal.activity.failing_activity": 2},
                 [True, False],
+                [1, 2],
                 id="activity_retry",
             ),
-            pytest.param(
-                ParentWorkflow,
-                40,
-                50,
-                {
-                    "temporal.workflow.ParentWorkflow": 1,
-                    "temporal.workflow.ChildWorkflow": 1,
-                    "temporal.activity.simple_activity": 1,
-                },
-                [False],
-                id="child_workflow",
-            ),
             # Local activities execute in the worker process and are traced like regular activities.
             pytest.param(
                 WorkflowWithLocalActivity,
                 100,
                 105,
                 {"temporal.workflow.WorkflowWithLocalActivity": 1, "temporal.activity.simple_local_activity": 1},
                 [False],
+                [1],
                 id="local_activity",
             ),
         ],
@@ -503,13 +477,14 @@ async def test_plugin_workflow_tracing(
         expected_result,
         expected_span_counts,
         expected_activity_errors,
+        expected_activity_attempts,
     ):
         task_queue = f"test-queue-{uuid.uuid4()}"
         async with Worker(
             temporal_env.client,
             task_queue=task_queue,
-            workflows=[WorkflowWithRetry, ParentWorkflow, ChildWorkflow, WorkflowWithLocalActivity],
-            activities=[failing_activity, simple_activity, simple_local_activity],
+            workflows=[WorkflowWithRetry, WorkflowWithLocalActivity],
+            activities=[failing_activity, simple_local_activity],
             plugins=[BraintrustPlugin(logger=memory_logger)],
         ):
             result = await temporal_env.client.execute_workflow(
@@ -528,6 +503,7 @@ async def test_plugin_workflow_tracing(
             key=lambda span: span["metrics"]["start"],
         )
         assert [bool(span.get("error")) for span in activity_spans] == expected_activity_errors
+        assert [span["metadata"]["temporal.activity_attempt"] for span in activity_spans] == expected_activity_attempts
         if any(expected_activity_errors):
             assert "Simulated failure on first attempt" in activity_spans[0]["error"]
 
@@ -546,22 +522,31 @@ async def test_child_workflow_and_activity_nest_under_in_workflow_span(
             workflow_task_executor=executor,
             plugins=[BraintrustPlugin(logger=memory_logger)],
         ):
-            await temporal_env.client.execute_workflow(
+            result = await temporal_env.client.execute_workflow(
                 DelegatingParentWorkflow.run,
                 TaskInput(value=1),
                 id=f"test-workflow-{uuid.uuid4()}",
                 task_queue=task_queue,
             )
+        assert result == 21
 
         braintrust.flush()
         spans = memory_logger.pop()
         # Every row of every span reaches the plugin's logger, whichever thread ran the activation.
         assert all("span_attributes" in span for span in spans)
         assert all(span["metrics"].get("end") is not None for span in spans)
+        assert Counter(span["span_attributes"]["name"] for span in spans) == {
+            "temporal.workflow.DelegatingParentWorkflow": 1,
+            "delegate_task": 1,
+            "temporal.workflow.ChildWorkflow": 1,
+            "temporal.activity.simple_activity": 2,
+        }
         (parent_wf,) = _spans_named(spans, "temporal.workflow.DelegatingParentWorkflow")
         (tool,) = _spans_named(spans, "delegate_task")
         (child_wf,) = _spans_named(spans, "temporal.workflow.ChildWorkflow")
         activities = _spans_named(spans, "temporal.activity.simple_activity")
+        assert all(not span.get("error") for span in activities)
+        assert all(span["metadata"]["temporal.activity_attempt"] == 1 for span in activities)
 
         assert tool["span_parents"] == [parent_wf["span_id"]]
         assert child_wf["span_parents"] == [tool["span_id"]]
@@ -571,6 +556,8 @@ async def test_child_workflow_and_activity_nest_under_in_workflow_span(
             if a["metadata"]["temporal.workflow_id"] != child_wf["metadata"]["temporal.workflow_id"]
         )
         assert parent_activity["span_parents"] == [tool["span_id"]]
+        (child_activity,) = [span for span in activities if span is not parent_activity]
+        assert child_activity["span_parents"] == [child_wf["span_id"]]
 
     @pytest.mark.asyncio
     async def test_in_workflow_spans_are_not_duplicated_on_replay(self, temporal_env, memory_logger):