diff --git a/py/src/braintrust/integrations/temporal/plugin.py b/py/src/braintrust/integrations/temporal/plugin.py index 2211244e..cb0c27e5 100644 --- a/py/src/braintrust/integrations/temporal/plugin.py +++ b/py/src/braintrust/integrations/temporal/plugin.py @@ -310,11 +310,7 @@ async def execute_workflow(self, input: temporalio.worker.ExecuteWorkflowInput) 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, - }, + metadata=_workflow_span_metadata(info), ) span.set_current() self._parent_span_context = _workflow_span_context(parent, ids) @@ -324,11 +320,7 @@ async def execute_workflow(self, input: temporalio.worker.ExecuteWorkflowInput) name=f"temporal.workflow.{info.workflow_type}", type="task", parent=parent_span_context or None, - metadata={ - "temporal.workflow_type": info.workflow_type, - "temporal.workflow_id": info.workflow_id, - "temporal.run_id": info.run_id, - }, + metadata=_workflow_span_metadata(info), ) span.set_current() self._parent_span_context = span.export() @@ -449,6 +441,18 @@ def _workflow_span_id_base(info: temporalio.workflow.Info) -> str: namespace = getattr(info, "namespace", "") return f"{_SPAN_ID_NAMESPACE}:{namespace}:{info.workflow_type}:{info.workflow_id}:{info.run_id}" +def _workflow_span_metadata(info: temporalio.workflow.Info) -> dict[str, Any]: + metadata: dict[str, Any] = { + "temporal.workflow_type": info.workflow_type, + "temporal.workflow_id": info.workflow_id, + "temporal.run_id": info.run_id, + } + # Added in temporalio 1.34.0. Unlike run_id, it stays the same when a workflow is reset. + original_execution_run_id = getattr(info, "original_execution_run_id", None) + if original_execution_run_id: + metadata["temporal.original_execution_run_id"] = original_execution_run_id + return metadata + def _workflow_span_ids(info: temporalio.workflow.Info, parent: str) -> dict[str, str]: return _deterministic_span_ids(_workflow_span_id_base(info), SpanComponentsV4.get_version(parent) < 4) diff --git a/py/src/braintrust/integrations/temporal/test_temporal.py b/py/src/braintrust/integrations/temporal/test_temporal.py index f6b6c6dd..a7fe13d5 100644 --- a/py/src/braintrust/integrations/temporal/test_temporal.py +++ b/py/src/braintrust/integrations/temporal/test_temporal.py @@ -2,6 +2,7 @@ import asyncio import concurrent.futures +import dataclasses import os import threading import uuid @@ -678,3 +679,34 @@ async def test_in_workflow_spans_are_not_duplicated_on_replay(self, temporal_env assert sorted(a["span_parents"][0] for a in activities) == sorted( [first_tool["span_id"], second_tool["span_id"]] ) + + @pytest.mark.asyncio + async def test_plugin_workflow_span_records_original_execution_run_id(self, temporal_env, memory_logger): + task_queue = f"test-queue-original-run-{uuid.uuid4()}" + async with Worker( + temporal_env.client, + task_queue=task_queue, + workflows=[WorkflowWithLocalActivity], + activities=[simple_local_activity], + plugins=[BraintrustPlugin(logger=memory_logger)], + ): + handle = await temporal_env.client.start_workflow( + WorkflowWithLocalActivity.run, + TaskInput(value=100), + id=f"test-workflow-original-run-{uuid.uuid4()}", + task_queue=task_queue, + ) + assert await handle.result() == 105 + + spans = memory_logger.pop() + workflow_spans = _spans_named(spans, "temporal.workflow.WorkflowWithLocalActivity") + assert len(workflow_spans) == 1 + metadata = workflow_spans[0]["metadata"] + assert metadata["temporal.run_id"] == handle.first_execution_run_id + + # temporalio 1.34.0 added Info.original_execution_run_id, which survives workflow resets + # (run_id changes on reset). Older pins don't have it, so the key is omitted there. + if "original_execution_run_id" in {field.name for field in dataclasses.fields(temporalio.workflow.Info)}: + assert metadata["temporal.original_execution_run_id"] == handle.first_execution_run_id + else: + assert "temporal.original_execution_run_id" not in metadata