💥 Move tracing uuids to use a separate Random from the workflow seed (#992)

* Move tracing uuids to use a separate Random from the workflow seed

* Fixing replay tests

* Switch to using a seed based on run_id

* Remove test failure

* Add comment

Author: tconley1428

temporalio/contrib/openai_agents/_temporal_trace_provider.py

@@ -15,6 +15,7 @@
 from agents.tracing.spans import Span
 
 from temporalio import workflow
+from temporalio.contrib.openai_agents._trace_interceptor import RunIdRandom
 from temporalio.workflow import ReadOnlyContextError
 
 
@@ -133,6 +134,13 @@ def force_flush(self) -> None:
         self._impl.force_flush()
 
 
+def _workflow_uuid() -> str:
+    random = cast(
+        RunIdRandom, getattr(workflow.instance(), "__temporal_openai_tracing_random")
+    )
+    return random.uuid4()
+
+
 class TemporalTraceProvider(DefaultTraceProvider):
     """A trace provider that integrates with Temporal workflows."""
 
@@ -156,7 +164,7 @@ def gen_trace_id(self) -> str:
         if workflow.in_workflow():
             try:
                 """Generate a new trace ID."""
-                return f"trace_{workflow.uuid4().hex}"
+                return f"trace_{_workflow_uuid()}"
             except ReadOnlyContextError:
                 return f"trace_{uuid.uuid4().hex}"
         return super().gen_trace_id()
@@ -166,7 +174,7 @@ def gen_span_id(self) -> str:
         if workflow.in_workflow():
             try:
                 """Generate a deterministic span ID."""
-                return f"span_{workflow.uuid4().hex[:24]}"
+                return f"span_{_workflow_uuid()}"
             except ReadOnlyContextError:
                 return f"span_{uuid.uuid4().hex[:24]}"
         return super().gen_span_id()
@@ -176,7 +184,7 @@ def gen_group_id(self) -> str:
         if workflow.in_workflow():
             try:
                 """Generate a deterministic group ID."""
-                return f"group_{workflow.uuid4().hex[:24]}"
+                return f"group_{_workflow_uuid()}"
             except ReadOnlyContextError:
                 return f"group_{uuid.uuid4().hex[:24]}"
         return super().gen_group_id()

temporalio/contrib/openai_agents/_trace_interceptor.py

@@ -2,16 +2,17 @@
 
 from __future__ import annotations
 
+import random
+import uuid
 from contextlib import contextmanager
-from typing import Any, Mapping, Protocol, Type, cast
+from typing import Any, Mapping, Protocol, Type
 
 from agents import CustomSpanData, custom_span, get_current_span, trace
 from agents.tracing import (
     get_trace_provider,
 )
-from agents.tracing.provider import DefaultTraceProvider
 from agents.tracing.scope import Scope
-from agents.tracing.spans import NoOpSpan, SpanImpl
+from agents.tracing.spans import NoOpSpan
 
 import temporalio.activity
 import temporalio.api.common.v1
@@ -283,6 +284,35 @@ async def execute_activity(
             return await self.next.execute_activity(input)
 
 
+class RunIdRandom:
+    """Random uuid generator seeded by the run id of the workflow.
+    Doesn't currently support replay over reset correctly.
+    """
+
+    def __init__(self):
+        """Create a new random UUID generator."""
+        self._random = random.Random("OpenAIPlugin" + workflow.info().run_id)
+
+    def uuid4(self) -> str:
+        """Generate a random UUID."""
+        return uuid.UUID(
+            bytes=random.getrandbits(16 * 8).to_bytes(16, "big"), version=4
+        ).hex[:24]
+
+
+def _ensure_tracing_random() -> None:
+    """We use a custom uuid generator for spans to ensure that changes to user code workflow.random usage
+    do not affect tracing and vice versa.
+    """
+    instance = workflow.instance()
+    if not hasattr(instance, "__temporal_openai_tracing_random"):
+        setattr(
+            workflow.instance(),
+            "__temporal_openai_tracing_random",
+            RunIdRandom(),
+        )
+
+
 class _ContextPropagationWorkflowInboundInterceptor(
     temporalio.worker.WorkflowInboundInterceptor
 ):
@@ -292,18 +322,21 @@ def init(self, outbound: temporalio.worker.WorkflowOutboundInterceptor) -> None:
     async def execute_workflow(
         self, input: temporalio.worker.ExecuteWorkflowInput
     ) -> Any:
+        _ensure_tracing_random()
         with context_from_header(
             "temporal:executeWorkflow", input, temporalio.workflow.payload_converter()
         ):
             return await self.next.execute_workflow(input)
 
     async def handle_signal(self, input: temporalio.worker.HandleSignalInput) -> None:
+        _ensure_tracing_random()
         with context_from_header(
             "temporal:handleSignal", input, temporalio.workflow.payload_converter()
         ):
             return await self.next.handle_signal(input)
 
     async def handle_query(self, input: temporalio.worker.HandleQueryInput) -> Any:
+        _ensure_tracing_random()
         with context_from_header(
             "temporal:handleQuery", input, temporalio.workflow.payload_converter()
         ):
@@ -322,6 +355,7 @@ def handle_update_validator(
     async def handle_update_handler(
         self, input: temporalio.worker.HandleUpdateInput
     ) -> Any:
+        _ensure_tracing_random()
         with context_from_header(
             "temporal:handleUpdateHandler",
             input,

tests/contrib/openai_agents/test_openai_replay.py

@@ -7,6 +7,9 @@
 from temporalio.contrib.openai_agents._temporal_openai_agents import (
     set_open_ai_agent_temporal_overrides,
 )
+from temporalio.contrib.openai_agents._trace_interceptor import (
+    OpenAIAgentsTracingInterceptor,
+)
 from temporalio.contrib.pydantic import pydantic_data_converter
 from temporalio.worker import Replayer
 from tests.contrib.openai_agents.test_openai import (
@@ -48,4 +51,5 @@ async def test_replay(file_name: str) -> None:
                     OutputGuardrailWorkflow,
                 ],
                 data_converter=pydantic_data_converter,
+                interceptors=[OpenAIAgentsTracingInterceptor()],
             ).replay_workflow(WorkflowHistory.from_json("fake", history_json))