Don't set startActivity as the parent context

Author: tconley1428

temporalio/contrib/openai_agents/_trace_interceptor.py

@@ -371,77 +371,60 @@ async def signal_child_workflow(
         self, input: temporalio.worker.SignalChildWorkflowInput
     ) -> None:
         trace = get_trace_provider().get_current_trace()
+        set_header_from_context(input, temporalio.workflow.payload_converter())
         if trace:
             with custom_span(
                 name="temporal:signalChildWorkflow",
                 data={"workflowId": input.child_workflow_id},
             ):
-                set_header_from_context(input, temporalio.workflow.payload_converter())
                 await self.next.signal_child_workflow(input)
-        else:
-            set_header_from_context(input, temporalio.workflow.payload_converter())
-            await self.next.signal_child_workflow(input)
+        await self.next.signal_child_workflow(input)
 
     async def signal_external_workflow(
         self, input: temporalio.worker.SignalExternalWorkflowInput
     ) -> None:
         trace = get_trace_provider().get_current_trace()
+        set_header_from_context(input, temporalio.workflow.payload_converter())
         if trace:
             with custom_span(
                 name="temporal:signalExternalWorkflow",
                 data={"workflowId": input.workflow_id},
             ):
-                set_header_from_context(input, temporalio.workflow.payload_converter())
                 await self.next.signal_external_workflow(input)
-        else:
-            set_header_from_context(input, temporalio.workflow.payload_converter())
-            await self.next.signal_external_workflow(input)
+        await self.next.signal_external_workflow(input)
 
     def start_activity(
         self, input: temporalio.worker.StartActivityInput
     ) -> temporalio.workflow.ActivityHandle:
         trace = get_trace_provider().get_current_trace()
-        span: Optional[Span] = None
+        set_header_from_context(input, temporalio.workflow.payload_converter())
         if trace:
-            span = custom_span(
+            with custom_span(
                 name="temporal:startActivity", data={"activity": input.activity}
-            )
-            span.start(mark_as_current=True)
-
-        set_header_from_context(input, temporalio.workflow.payload_converter())
-        handle = self.next.start_activity(input)
-        if span:
-            handle.add_done_callback(lambda _: span.finish())  # type: ignore
-        return handle
+            ):
+                return self.next.start_activity(input)
+        return self.next.start_activity(input)
 
     async def start_child_workflow(
         self, input: temporalio.worker.StartChildWorkflowInput
     ) -> temporalio.workflow.ChildWorkflowHandle:
         trace = get_trace_provider().get_current_trace()
-        span: Optional[Span] = None
+        set_header_from_context(input, temporalio.workflow.payload_converter())
         if trace:
-            span = custom_span(
+            with custom_span(
                 name="temporal:startChildWorkflow", data={"workflow": input.workflow}
-            )
-            span.start(mark_as_current=True)
-        set_header_from_context(input, temporalio.workflow.payload_converter())
-        handle = await self.next.start_child_workflow(input)
-        if span:
-            handle.add_done_callback(lambda _: span.finish())  # type: ignore
-        return handle
+            ):
+                return await self.next.start_child_workflow(input)
+        return await self.next.start_child_workflow(input)
 
     def start_local_activity(
         self, input: temporalio.worker.StartLocalActivityInput
     ) -> temporalio.workflow.ActivityHandle:
         trace = get_trace_provider().get_current_trace()
-        span: Optional[Span] = None
+        set_header_from_context(input, temporalio.workflow.payload_converter())
         if trace:
-            span = custom_span(
+            with custom_span(
                 name="temporal:startLocalActivity", data={"activity": input.activity}
-            )
-            span.start(mark_as_current=True)
-        set_header_from_context(input, temporalio.workflow.payload_converter())
-        handle = self.next.start_local_activity(input)
-        if span:
-            handle.add_done_callback(lambda _: span.finish())  # type: ignore
-        return handle
+            ):
+                return self.next.start_local_activity(input)
+        return self.next.start_local_activity(input)

tests/contrib/openai_agents/test_openai_tracing.py

@@ -81,18 +81,22 @@ def paired_span(a: tuple[Span[Any], bool], b: tuple[Span[Any], bool]) -> None:
     # Initial planner spans - There are only 3 because we don't make an actual model call
     paired_span(processor.span_events[0], processor.span_events[5])
     assert processor.span_events[0][0].span_data.export().get("name") == "PlannerAgent"
+    print("Planner span:", processor.span_events[0][0].span_id)
 
-    paired_span(processor.span_events[1], processor.span_events[4])
+    # startActivity happens before executeActivity
+    paired_span(processor.span_events[1], processor.span_events[2])
     assert (
         processor.span_events[1][0].span_data.export().get("name")
         == "temporal:startActivity"
     )
+    print("StartActivity span:", processor.span_events[1][0].span_id)
 
-    paired_span(processor.span_events[2], processor.span_events[3])
+    paired_span(processor.span_events[3], processor.span_events[4])
     assert (
-        processor.span_events[2][0].span_data.export().get("name")
+        processor.span_events[3][0].span_data.export().get("name")
         == "temporal:executeActivity"
     )
+    print("Execute Span parent:", processor.span_events[3][0].parent_id)
 
     for span, start in processor.span_events[6:-6]:
         span_data = span.span_data.export()
@@ -113,28 +117,27 @@ def paired_span(a: tuple[Span[Any], bool], b: tuple[Span[Any], bool]) -> None:
                 len(parents) == 2 and parents[0].span_data.export()["type"] == "agent"
             )
 
-        # Execute is parented to start
+        # Execute is parented to the agent as well
         if span_data.get("name") == "temporal:executeActivity":
             parents = [
                 s for (s, _) in processor.span_events if s.span_id == span.parent_id
             ]
             assert (
-                len(parents) == 2
-                and parents[0].span_data.export()["name"] == "temporal:startActivity"
+                len(parents) == 2 and parents[0].span_data.export()["type"] == "agent"
             )
 
     # Final writer spans - There are only 3 because we don't make an actual model call
     paired_span(processor.span_events[-6], processor.span_events[-1])
     assert processor.span_events[-6][0].span_data.export().get("name") == "WriterAgent"
 
-    paired_span(processor.span_events[-5], processor.span_events[-2])
+    paired_span(processor.span_events[-5], processor.span_events[-4])
     assert (
         processor.span_events[-5][0].span_data.export().get("name")
         == "temporal:startActivity"
     )
 
-    paired_span(processor.span_events[-4], processor.span_events[-3])
+    paired_span(processor.span_events[-3], processor.span_events[-2])
     assert (
-        processor.span_events[-4][0].span_data.export().get("name")
+        processor.span_events[-3][0].span_data.export().get("name")
         == "temporal:executeActivity"
     )