Fix Temporal replay span parent propagation

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

@@ -80,6 +80,8 @@
 """
 
 import dataclasses
+import hashlib
+import uuid
 from collections.abc import Mapping
 from typing import Any
 
@@ -90,6 +92,8 @@
 import temporalio.converter
 import temporalio.worker
 import temporalio.workflow
+from braintrust.span_identifier_v3 import SpanComponentsV3
+from braintrust.span_identifier_v4 import SpanComponentsV4
 from temporalio.plugin import SimplePlugin
 from temporalio.worker import WorkflowRunner
 from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner
@@ -112,6 +116,8 @@
 
 # Header key for passing span context between client, workflows, and activities
 _HEADER_KEY = "_braintrust-span"
+_SPAN_ID_NAMESPACE = "braintrust.temporal.workflow"
+_SPAN_CONTEXT_PATCH_ID = "braintrust-temporal-workflow-span-context-v1"
 
 
 class BraintrustInterceptor(temporalio.client.Interceptor, temporalio.worker.Interceptor):
@@ -173,7 +179,7 @@ def workflow_interceptor_class(
 
     def _span_context_to_headers(
         self,
-        span_context: dict[str, Any],
+        span_context: Any,
         headers: Mapping[str, temporalio.api.common.v1.Payload],
     ) -> Mapping[str, temporalio.api.common.v1.Payload]:
         """Add span context to headers."""
@@ -186,9 +192,7 @@ def _span_context_to_headers(
                 }
         return headers
 
-    def _span_context_from_headers(
-        self, headers: Mapping[str, temporalio.api.common.v1.Payload]
-    ) -> dict[str, Any] | None:
+    def _span_context_from_headers(self, headers: Mapping[str, temporalio.api.common.v1.Payload]) -> Any | None:
         """Extract span context from headers."""
         if _HEADER_KEY not in headers:
             return None
@@ -277,7 +281,8 @@ class BraintrustWorkflowInboundInterceptor(temporalio.worker.WorkflowInboundInte
     def __init__(self, next: temporalio.worker.WorkflowInboundInterceptor) -> None:
         super().__init__(next)
         self.payload_converter = temporalio.converter.PayloadConverter.default
-        self._parent_span_context: dict[str, Any] | None = None
+        self._parent_span_context: Any | None = None
+        self._workflow_span: Any | None = None
 
     def init(self, outbound: temporalio.worker.WorkflowOutboundInterceptor) -> None:
         super().init(_BraintrustWorkflowOutboundInterceptor(outbound, self))
@@ -298,14 +303,35 @@ async def execute_workflow(self, input: temporalio.worker.ExecuteWorkflowInput)
         # Create a span for the workflow execution using sandbox_unrestricted
         # to bypass the sandbox restrictions on logger state access
         span = None
-        if not temporalio.workflow.unsafe.is_replaying():
-            with temporalio.workflow.unsafe.sandbox_unrestricted():
-                # Get logger via extern function (supports test logger parameter)
-                get_logger = temporalio.workflow.extern_functions()["__braintrust_get_logger"]
-                logger = get_logger()
-
-                if logger:
-                    info = temporalio.workflow.info()
+        with temporalio.workflow.unsafe.sandbox_unrestricted():
+            # Get logger via extern function (supports test logger parameter)
+            get_logger = temporalio.workflow.extern_functions()["__braintrust_get_logger"]
+            logger = get_logger()
+
+            if logger:
+                info = temporalio.workflow.info()
+                parent = parent_span_context or logger.export()
+                if temporalio.workflow.patched(_SPAN_CONTEXT_PATCH_ID):
+                    ids = _workflow_span_ids(info, parent)
+                    if temporalio.workflow.unsafe.is_replaying():
+                        self._parent_span_context = _workflow_span_context(parent, ids)
+                    else:
+                        span = logger.start_span(
+                            name=f"temporal.workflow.{info.workflow_type}",
+                            type="task",
+                            parent=parent,
+                            id=ids["row_id"],
+                            span_id=ids["span_id"],
+                            root_span_id=ids["root_span_id"],
+                            metadata={
+                                "temporal.workflow_type": info.workflow_type,
+                                "temporal.workflow_id": info.workflow_id,
+                                "temporal.run_id": info.run_id,
+                            },
+                        )
+                        span.set_current()
+                        self._parent_span_context = _workflow_span_context(parent, ids)
+                elif not temporalio.workflow.unsafe.is_replaying():
                     span = logger.start_span(
                         name=f"temporal.workflow.{info.workflow_type}",
                         type="task",
@@ -317,23 +343,23 @@ async def execute_workflow(self, input: temporalio.worker.ExecuteWorkflowInput)
                         },
                     )
                     span.set_current()
-
-                    # Update parent span context for activities
                     self._parent_span_context = span.export()
 
+        self._workflow_span = span
+
         try:
             result = await super().execute_workflow(input)
             return result
         except Exception as e:
-            if span:
+            if self._workflow_span:
                 with temporalio.workflow.unsafe.sandbox_unrestricted():
-                    span.log(error=str(e))
+                    self._workflow_span.log(error=str(e))
             raise
         finally:
-            if span:
+            if self._workflow_span:
                 with temporalio.workflow.unsafe.sandbox_unrestricted():
-                    span.unset_current()
-                    span.end()
+                    self._workflow_span.unset_current()
+                    self._workflow_span.end()
 
 
 class _BraintrustWorkflowOutboundInterceptor(temporalio.worker.WorkflowOutboundInterceptor):
@@ -388,6 +414,37 @@ def _modify_workflow_runner(existing: WorkflowRunner | None) -> WorkflowRunner |
     return existing
 
 
+def _workflow_span_ids(info: temporalio.workflow.Info, parent: str) -> dict[str, str]:
+    namespace = getattr(info, "namespace", "")
+    base = f"{_SPAN_ID_NAMESPACE}:{namespace}:{info.workflow_type}:{info.workflow_id}:{info.run_id}"
+    if SpanComponentsV4.get_version(parent) >= 4:
+        digest = hashlib.sha256(base.encode()).hexdigest()
+        return {
+            "row_id": str(uuid.uuid5(uuid.NAMESPACE_URL, base)),
+            "span_id": digest[:16],
+            "root_span_id": digest[16:48],
+        }
+    return {
+        "row_id": str(uuid.uuid5(uuid.NAMESPACE_URL, f"{base}:row")),
+        "span_id": str(uuid.uuid5(uuid.NAMESPACE_URL, f"{base}:span")),
+        "root_span_id": str(uuid.uuid5(uuid.NAMESPACE_URL, f"{base}:root")),
+    }
+
+
+def _workflow_span_context(parent: str, ids: dict[str, str]) -> str:
+    component_cls = SpanComponentsV4 if SpanComponentsV4.get_version(parent) >= 4 else SpanComponentsV3
+    parent_components = component_cls.from_str(parent)
+    return component_cls(
+        object_type=parent_components.object_type,
+        object_id=parent_components.object_id,
+        compute_object_metadata_args=parent_components.compute_object_metadata_args,
+        row_id=ids["row_id"],
+        span_id=ids["span_id"],
+        root_span_id=parent_components.root_span_id or ids["root_span_id"],
+        propagated_event=parent_components.propagated_event,
+    ).to_str()
+
+
 class BraintrustPlugin(SimplePlugin):
     """Braintrust plugin for Temporal that automatically configures tracing.
 

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

@@ -5,7 +5,7 @@
 import uuid
 from dataclasses import dataclass
 from datetime import timedelta
-from typing import Any
+from typing import Any, cast
 
 import pytest
 import pytest_asyncio
@@ -22,11 +22,23 @@
 import temporalio.worker
 import temporalio.workflow
 from braintrust.integrations.temporal import BraintrustInterceptor, BraintrustPlugin
-from braintrust.test_helpers import init_test_logger
+from braintrust.integrations.temporal.plugin import _workflow_span_context, _workflow_span_ids
+from braintrust.span_identifier_v3 import SpanComponentsV3, SpanObjectTypeV3
+from braintrust.span_identifier_v4 import SpanComponentsV4
+from braintrust.test_helpers import init_test_logger, preserve_env_vars
+from temporalio.client import Client
 from temporalio.common import RetryPolicy
 from temporalio.worker import Worker
 
 
+@dataclass
+class WorkflowInfoForTest:
+    namespace: str
+    workflow_type: str
+    workflow_id: str
+    run_id: str
+
+
 class TestHeaderSerialization:
     """Unit tests for header serialization/deserialization."""
 
@@ -103,6 +115,56 @@ def test_span_context_roundtrip(self):
         assert result_context == original_context
 
 
+class TestWorkflowSpanContext:
+    def test_workflow_span_context_preserves_legacy_parent_encoding(self):
+        parent_components = SpanComponentsV3(
+            object_type=SpanObjectTypeV3.PROJECT_LOGS,
+            object_id=str(uuid.uuid4()),
+            row_id=str(uuid.uuid4()),
+            span_id=str(uuid.uuid4()),
+            root_span_id=str(uuid.uuid4()),
+        )
+        parent = parent_components.to_str()
+        info = WorkflowInfoForTest(
+            namespace="default",
+            workflow_type="ReplayAfterSignalWorkflow",
+            workflow_id="workflow-id",
+            run_id="run-id",
+        )
+
+        with preserve_env_vars("BRAINTRUST_LEGACY_IDS"):
+            os.environ.pop("BRAINTRUST_LEGACY_IDS", None)
+            ids = _workflow_span_ids(cast(temporalio.workflow.Info, info), parent)
+            context = _workflow_span_context(parent, ids)
+
+        assert SpanComponentsV4.get_version(context) == 3
+        parsed = SpanComponentsV3.from_str(context)
+        assert parsed.row_id == ids["row_id"]
+        assert parsed.span_id == ids["span_id"]
+        assert parsed.root_span_id == parent_components.root_span_id
+
+    def test_workflow_span_context_uses_stable_root_for_object_parent(self):
+        parent = SpanComponentsV4(
+            object_type=SpanObjectTypeV3.PROJECT_LOGS,
+            object_id=str(uuid.uuid4()),
+        ).to_str()
+        info = WorkflowInfoForTest(
+            namespace="default",
+            workflow_type="ReplayAfterSignalWorkflow",
+            workflow_id="workflow-id",
+            run_id="run-id",
+        )
+
+        ids = _workflow_span_ids(cast(temporalio.workflow.Info, info), parent)
+        context = _workflow_span_context(parent, ids)
+
+        assert SpanComponentsV4.get_version(context) == 4
+        parsed = SpanComponentsV4.from_str(context)
+        assert parsed.row_id == ids["row_id"]
+        assert parsed.span_id == ids["span_id"]
+        assert parsed.root_span_id == ids["root_span_id"]
+
+
 # Integration Test Infrastructure
 
 
@@ -222,6 +284,39 @@ async def run(self, input: TaskInput) -> int:
         return child_result
 
 
+@temporalio.workflow.defn
+class ReplayAfterSignalWorkflow:
+    """Workflow that schedules an activity after a later workflow task."""
+
+    def __init__(self) -> None:
+        self._continue = False
+        self._state = "starting"
+
+    @temporalio.workflow.run
+    async def run(self, input: TaskInput) -> int:
+        first_result = await temporalio.workflow.execute_activity(
+            simple_activity,
+            input,
+            start_to_close_timeout=timedelta(seconds=10),
+        )
+        self._state = "waiting"
+        await temporalio.workflow.wait_condition(lambda: self._continue)
+        self._state = "continued"
+        return await temporalio.workflow.execute_activity(
+            simple_activity,
+            TaskInput(value=first_result),
+            start_to_close_timeout=timedelta(seconds=10),
+        )
+
+    @temporalio.workflow.signal
+    def continue_workflow(self) -> None:
+        self._continue = True
+
+    @temporalio.workflow.query
+    def state(self) -> str:
+        return self._state
+
+
 class TestAutoInstrumentation:
     """Tests for Temporal auto-instrumentation helpers."""
 
@@ -271,6 +366,18 @@ def memory_logger():
         yield bgl
 
 
+def _spans_named(spans: list[dict[str, Any]], name: str) -> list[dict[str, Any]]:
+    return [span for span in spans if span.get("span_attributes", {}).get("name") == name]
+
+
+async def _wait_for_workflow_state(handle: Any, expected: str) -> None:
+    for _ in range(50):
+        if await handle.query(ReplayAfterSignalWorkflow.state) == expected:
+            return
+        await asyncio.sleep(0.1)
+    raise AssertionError(f"Workflow did not reach state {expected!r}")
+
+
 class TestBraintrustPluginIntegration:
     """Integration tests for BraintrustPlugin with real Temporal workflows."""
 
@@ -369,6 +476,63 @@ async def test_plugin_context_propagation(self, temporal_env, memory_logger):
         assert len(workflow_spans) > 0, "Expected workflow spans"
         assert len(activity_spans) > 0, "Expected activity spans"
 
+    @pytest.mark.parametrize("with_client_parent", [False, True])
+    @pytest.mark.asyncio
+    async def test_plugin_activity_after_replay_stays_under_workflow_span(
+        self, temporal_env, memory_logger, with_client_parent
+    ):
+        task_queue = f"test-queue-replay-{with_client_parent}-{uuid.uuid4()}"
+        workflow_id = f"test-workflow-replay-{with_client_parent}-{uuid.uuid4()}"
+        workflow_client = temporal_env.client
+        if with_client_parent:
+            workflow_client = await Client.connect(
+                temporal_env.client.service_client.config.target_host,
+                namespace=temporal_env.client.namespace,
+                plugins=[BraintrustPlugin(logger=memory_logger)],
+            )
+
+        async with Worker(
+            temporal_env.client,
+            task_queue=task_queue,
+            workflows=[ReplayAfterSignalWorkflow],
+            activities=[simple_activity],
+            max_cached_workflows=0,
+            plugins=[BraintrustPlugin(logger=memory_logger)],
+        ):
+            if with_client_parent:
+                with braintrust.start_span(name="test.client_replay_operation", type="task"):
+                    handle = await workflow_client.start_workflow(
+                        ReplayAfterSignalWorkflow.run,
+                        TaskInput(value=10),
+                        id=workflow_id,
+                        task_queue=task_queue,
+                    )
+            else:
+                handle = await workflow_client.start_workflow(
+                    ReplayAfterSignalWorkflow.run,
+                    TaskInput(value=10),
+                    id=workflow_id,
+                    task_queue=task_queue,
+                )
+            await _wait_for_workflow_state(handle, "waiting")
+            await handle.signal(ReplayAfterSignalWorkflow.continue_workflow)
+            assert await handle.result() == 30
+
+        braintrust.flush()
+        spans = memory_logger.pop()
+        workflow_spans = _spans_named(spans, "temporal.workflow.ReplayAfterSignalWorkflow")
+        activity_spans = _spans_named(spans, "temporal.activity.simple_activity")
+
+        assert len(workflow_spans) == 1
+        assert len(activity_spans) == 2
+        workflow_span = workflow_spans[0]
+        assert all(workflow_span["span_id"] in span.get("span_parents", []) for span in activity_spans)
+        assert all(workflow_span["root_span_id"] == span["root_span_id"] for span in activity_spans)
+        if with_client_parent:
+            client_spans = _spans_named(spans, "test.client_replay_operation")
+            assert len(client_spans) == 1
+            assert workflow_span["root_span_id"] == client_spans[0]["root_span_id"]
+
     @pytest.mark.asyncio
     async def test_plugin_activity_retry_tracing(self, temporal_env, memory_logger):
         """Test that activity retries are properly traced.