From 3cbbc92f54bdf35f4f8ad9fe81a754b877815790 Mon Sep 17 00:00:00 2001 From: Thibault Soubeste Date: Thu, 9 Jul 2026 09:26:37 +0900 Subject: [PATCH] fix: keep Temporal replay spans under workflow --- .../integrations/temporal/plugin.py | 97 +++++++--- .../integrations/temporal/test_temporal.py | 168 +++++++++++++++++- 2 files changed, 243 insertions(+), 22 deletions(-) diff --git a/py/src/braintrust/integrations/temporal/plugin.py b/py/src/braintrust/integrations/temporal/plugin.py index b25a3946..669a9fe0 100644 --- a/py/src/braintrust/integrations/temporal/plugin.py +++ b/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. diff --git a/py/src/braintrust/integrations/temporal/test_temporal.py b/py/src/braintrust/integrations/temporal/test_temporal.py index 671ffd1e..da2c4bba 100644 --- a/py/src/braintrust/integrations/temporal/test_temporal.py +++ b/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.