diff --git a/CHANGELOG.md b/CHANGELOG.md index f1a526e7e..4879440ae 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/temporalio/contrib/openai_agents/_trace_interceptor.py b/temporalio/contrib/openai_agents/_trace_interceptor.py index 66297e20b..b262ce5da 100644 --- a/temporalio/contrib/openai_agents/_trace_interceptor.py +++ b/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 diff --git a/tests/contrib/openai_agents/test_openai_tracing.py b/tests/contrib/openai_agents/test_openai_tracing.py index 28b804cc1..88924dd5e 100644 --- a/tests/contrib/openai_agents/test_openai_tracing.py +++ b/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"