Fix openai_agents tracing leaving temporal:startActivity as the current span

CHANGELOG.md

@@ -86,6 +86,11 @@ to include examples, links to docs, or any other relevant information.
 - Nexus-context workflow/activity starts no longer set `on_conflict_options` when there are no links
   or callbacks to attach.
 - The workflow sandbox now passes `pydantic_core` through by default, alongside `pydantic`.
+- `temporalio.contrib.openai_agents`: the `temporal:startActivity`, `temporal:startChildWorkflow`,
+  and `temporal:startLocalActivity` spans no longer remain the current Agents SDK span after the
+  start. Previously, spans created afterwards in the same context were parented to them; in an
+  agent loop, tool calls that followed a model call nested under the model call's span instead
+  of under the turn.
 
 ### Security
 

temporalio/contrib/openai_agents/_trace_interceptor.py

@@ -60,6 +60,24 @@ def temporal_span(
         yield
 
 
+@contextmanager
+def _as_current_span(span: Span[Any] | None):
+    """Make ``span`` the current span for the duration of the block only.
+
+    Used to capture a span in an outbound header without leaving it current,
+    which would make it the parent of spans created later in the same context,
+    such as the tool calls that follow a model call within an agent turn.
+    """
+    if span is None:
+        yield
+        return
+    token = Scope.set_current_span(span)
+    try:
+        yield
+    finally:
+        Scope.reset_current_span(token)
+
+
 class OpenAIAgentsContextPropagationInterceptor(
     temporalio.client.Interceptor, temporalio.worker.Interceptor
 ):
@@ -404,9 +422,9 @@ def start_activity(
             span = custom_span(
                 name="temporal:startActivity", data={"activity": input.activity}
             )
-            span.start(mark_as_current=True)
-
-        self.root().set_header_from_context(input)
+            span.start()
+        with _as_current_span(span):
+            self.root().set_header_from_context(input)
         handle = self.next.start_activity(input)
         if span:
             handle.add_done_callback(lambda _: span.finish())  # type: ignore
@@ -421,8 +439,9 @@ async def start_child_workflow(
             span = custom_span(
                 name="temporal:startChildWorkflow", data={"workflow": input.workflow}
             )
-            span.start(mark_as_current=True)
-        self.root().set_header_from_context(input)
+            span.start()
+        with _as_current_span(span):
+            self.root().set_header_from_context(input)
         handle = await self.next.start_child_workflow(input)
         if span:
             handle.add_done_callback(lambda _: span.finish())  # type: ignore
@@ -437,8 +456,9 @@ def start_local_activity(
             span = custom_span(
                 name="temporal:startLocalActivity", data={"activity": input.activity}
             )
-            span.start(mark_as_current=True)
-        self.root().set_header_from_context(input)
+            span.start()
+        with _as_current_span(span):
+            self.root().set_header_from_context(input)
         handle = self.next.start_local_activity(input)
         if span:
             handle.add_done_callback(lambda _: span.finish())  # type: ignore

tests/contrib/openai_agents/test_openai_tracing.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import opentelemetry.trace
-from agents import Span, Trace, TracingProcessor, custom_span, trace
+from agents import Agent, Runner, Span, Trace, TracingProcessor, custom_span, trace
 from agents.tracing import get_trace_provider
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
@@ -14,7 +14,10 @@
 from temporalio.contrib.openai_agents import _temporal_openai_agents
 from temporalio.contrib.openai_agents.testing import (
     AgentEnvironment,
+    ResponseBuilders,
+    TestModel,
 )
+from temporalio.contrib.openai_agents.workflow import activity_as_tool
 from temporalio.contrib.opentelemetry import create_tracer_provider
 from temporalio.worker.workflow_sandbox import (
     SandboxedWorkflowRunner,
@@ -237,6 +240,78 @@ def paired_span(a: tuple[Span[Any], bool], b: tuple[Span[Any], bool]) -> None:
         )
 
 
+@activity.defn
+async def lookup_account(account_id: str) -> str:
+    return f"account {account_id}"
+
+
+@workflow.defn
+class ToolTracingWorkflow:
+    @workflow.run
+    async def run(self) -> str:
+        agent = Agent[str](
+            name="Account agent",
+            instructions="Look up the account.",
+            tools=[
+                activity_as_tool(
+                    lookup_account, start_to_close_timeout=timedelta(seconds=10)
+                )
+            ],
+        )
+        result = await Runner.run(agent, "Look up account 1")
+        return result.final_output
+
+
+async def test_tool_span_parented_to_turn(client: Client):
+    """A tool call that follows a model call in the same turn is a sibling of the
+    model call's temporal:startActivity span, not its child."""
+    model = TestModel.returning_responses(
+        [
+            ResponseBuilders.tool_call('{"account_id": "1"}', "lookup_account"),
+            ResponseBuilders.output_message("done"),
+        ]
+    )
+    async with AgentEnvironment(model=model) as env:
+        client = env.applied_on_client(client)
+        processor = MemoryTracingProcessor()
+        get_trace_provider().set_processors([processor])
+
+        async with new_worker(
+            client, ToolTracingWorkflow, activities=[lookup_account]
+        ) as worker:
+            with trace("Tool workflow") as t:
+                await client.execute_workflow(
+                    ToolTracingWorkflow.run,
+                    id=f"tool-tracing-workflow-{uuid.uuid4()}",
+                    task_queue=worker.task_queue,
+                    execution_timeout=timedelta(seconds=120),
+                )
+
+    # MemoryTracingProcessor's lists are shared across tests; keep only this trace
+    spans = {
+        s.span_id: s
+        for s, started in processor.span_events
+        if started and s.trace_id == t.trace_id
+    }
+
+    def name(span: Span[Any]) -> str | None:
+        return span.span_data.export().get("name")
+
+    def parent_name(span: Span[Any]) -> str | None:
+        return name(spans[span.parent_id]) if span.parent_id else None
+
+    tool_span = next(s for s in spans.values() if s.span_data.type == "function")
+    assert parent_name(tool_span) == "turn"
+
+    # Model calls (one per turn) stay under their turn; the tool's activity stays under the tool
+    start_spans = [s for s in spans.values() if name(s) == "temporal:startActivity"]
+    assert sorted(parent_name(s) or "" for s in start_spans) == [
+        "lookup_account",
+        "turn",
+        "turn",
+    ]
+
+
 @activity.defn
 async def simple_no_context_activity() -> str:
     return "success"